amend

Author: vmoens

test/test_collector.py

@@ -3938,13 +3938,12 @@ def _maybe_map_weights(self, server_weights: TensorDictBase) -> TensorDictBase:
         def all_worker_ids(self) -> list[int] | list[torch.device]:
             return list(range(self.num_workers))
 
-    @pytest.mark.skipif(not _has_cuda, reason="requires cuda another device than CPU.")
     @pytest.mark.skipif(not _has_gym, reason="requires gym")
     @pytest.mark.parametrize(
         "weight_updater", ["scheme_shared", "scheme_pipe", "weight_updater"]
     )
     def test_weight_update(self, weight_updater):
-        device = "cuda:0"
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
         env_maker = lambda: GymEnv(PENDULUM_VERSIONED(), device="cpu")
         policy_factory = lambda: TensorDictModule(
             nn.Linear(3, 1, device=device), in_keys=["observation"], out_keys=["action"]

torchrl/collectors/_multi_base.py

@@ -835,9 +835,9 @@ def _run_processes(self) -> None:
         # can be initialized here since all required resources exist
         if self._weight_sync_schemes:
             for model_id, scheme in self._weight_sync_schemes.items():
-                if hasattr(scheme, "init_on_sender"):
+                if not scheme.initialized_on_sender:
                     scheme.init_on_sender(model_id=model_id, context=self)
-                    self._weight_senders[model_id] = scheme.get_sender()
+                self._weight_senders[model_id] = scheme.get_sender()
 
         # Create a policy on the right device
         policy_factory = self.policy_factory

torchrl/collectors/_runner.py

@@ -40,7 +40,9 @@ def _make_policy_factory(
     if weight_sync_scheme is not None:
         # Initialize the receiver on the worker side
         weight_sync_scheme.init_on_receiver(
-            model=policy, model_id="policy", worker_idx=worker_idx, pipe=pipe
+            model=policy,
+            model_id="policy",
+            worker_idx=worker_idx,
         )
         # Get the receiver and synchronize initial weights
         receiver = weight_sync_scheme.get_receiver()