amend

pytorch · vmoens · Nov 30, 2023 · Oct 6, 2023 · Oct 6, 2023 · Oct 6, 2023
commit 939ece493da01ee00715d19456839278183fbf72
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -1438,8 +1438,8 @@ def _run_worker_pipe_cuda(
                     env_input = shared_tensordict
                 td, root_next_td = env.step_and_maybe_reset(env_input)
                 if env_device_cpu:
-                    next_shared_tensordict.apply(_update_cuda, td.get("next"))
-                    shared_tensordict.apply(_update_cuda, root_next_td)
+                    next_shared_tensordict._fast_apply(_update_cuda, td.get("next"))
+                    shared_tensordict._fast_apply(_update_cuda, root_next_td)
                 else:
                     next_shared_tensordict.update_(td.get("next"))
                     shared_tensordict.update_(root_next_td)