Merge remote-tracking branch 'origin/version-0.5' into version-0.5

torchrl/envs/batched_envs.py

@@ -406,17 +406,16 @@ def _find_sync_values(self):
             return _do_nothing, _do_nothing
 
         if worker_device is None:
-            worker_not_main = [False]
+            worker_not_main = False
 
-            def find_all_worker_devices(item, worker_not_main=worker_not_main):
+            def find_all_worker_devices(item):
+                nonlocal worker_not_main
                 if hasattr(item, "device"):
-                    worker_not_main[0] = worker_not_main[0] or (
-                        item.device != self_device
-                    )
+                    worker_not_main = worker_not_main or (item.device != self_device)
 
             for td in self.shared_tensordicts:
                 td.apply(find_all_worker_devices, filter_empty=True)
-            if worker_not_main[0]:
+            if worker_not_main:
                 if torch.cuda.is_available():
                     worker_device = (
                         torch.device("cuda")
@@ -431,6 +430,8 @@ def find_all_worker_devices(item, worker_not_main=worker_not_main):
                     )
                 else:
                     raise RuntimeError("Did not find a valid worker device")
+            else:
+                worker_device = self_device
 
         if (
             worker_device is not None
@@ -460,6 +461,7 @@ def find_all_worker_devices(item, worker_not_main=worker_not_main):
             and self_device.type == "mps"
         ):
             return _mps_sync(self_device), _mps_sync(self_device)
+        return _do_nothing, _do_nothing
 
     def __getstate__(self):
         out = copy(self.__dict__)