Update

[ghstack-poisoned]

Author: vmoens

test/test_env.py

@@ -4142,12 +4142,13 @@ def test_parallel_partial_step_and_maybe_reset(
             )
             try:
                 td = penv.reset()
-                psteps = torch.zeros(4, dtype=torch.bool)
+                psteps = torch.zeros(4, dtype=torch.bool, device=td.get("done").device)
                 psteps[[1, 3]] = True
                 td.set("_step", psteps)
 
                 td.set("action", penv.full_action_spec[penv.action_key].one())
                 td, tdreset = penv.step_and_maybe_reset(td)
+                print(td)
                 assert_allclose_td(td[0].get("next"), td[0], intersection=True)
                 assert (td[1].get("next") != 0).any()
                 assert_allclose_td(td[2].get("next"), td[2], intersection=True)