Change dtype for seq_lens to int32

Signed-off-by: Xuehai Pan <[email protected]>
ray-project · Apr 27, 2022 · 426ff69 · 426ff69
1 parent 8b6738e
commit 426ff69
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/rllib/policy/tests/test_rnn_sequencing.py b/rllib/policy/tests/test_rnn_sequencing.py
@@ -100,7 +100,11 @@ def test_pad_batch_fixed_max(self):
     def test_add_time_dimension(self):
         """Test add_time_dimension gives sequential data along the time dimension"""
 
-        B, T, F = np.random.choice(list(range(8, 32)), size=3, replace=False)
+        B, T, F = np.random.choice(
+            np.asarray(list(range(8, 32)), dtype=np.int32),  # use int32 for seq_lens
+            size=3,
+            replace=False,
+        )
 
         inputs_numpy = np.repeat(
             np.arange(B * T)[:, np.newaxis], repeats=F, axis=-1