Fix TF Lite unsupported operation

m-wojnar · Jul 15, 2023 · bdc4d0c · bdc4d0c
1 parent 7a423aa
commit bdc4d0c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/reinforced_lib/utils/experience_replay.py b/reinforced_lib/utils/experience_replay.py
@@ -166,7 +166,7 @@ def sample(buffer: ReplayBuffer, key: PRNGKey) -> Tuple:
             Tuple containing the batch of states, actions, rewards, terminals and next states.
         """
 
-        idxs = jax.random.randint(key, shape=(batch_size,), minval=0, maxval=buffer.size)
+        idxs = jax.random.uniform(key, shape=(batch_size,), minval=0, maxval=buffer.size).astype(jnp.int32)
 
         states = buffer.states[idxs]
         actions = buffer.actions[idxs]