Update

[ghstack-poisoned]

Author: vmoens

test/test_env.py

@@ -4148,7 +4148,6 @@ def test_parallel_partial_step_and_maybe_reset(
 
                 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)

torchrl/envs/batched_envs.py

@@ -1691,6 +1691,14 @@ def step_and_maybe_reset(
                 device=device,
                 filter_empty=True,
             )
+            if tensordict.device != device:
+                tensordict = tensordict._fast_apply(
+                    lambda x: x.to(device, non_blocking=self.non_blocking)
+                    if x.device != device
+                    else x,
+                    device=device,
+                    filter_empty=True,
+                )
             self._sync_w2m()
         else:
             next_td = next_td.clone().clear_device_()