[BugFix] load_state_dict in param updates for collectors (#1145)

pytorch · May 10, 2023 · a48cd95 · a48cd95
1 parent 9402664
commit a48cd95
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -300,7 +300,7 @@ def update_policy_weights_(
         if policy_weights is not None:
             self.policy_weights.apply(lambda x: x.data).update_(policy_weights)
         elif self.get_weights_fn is not None:
-            self.policy.load_state_dict(self.get_weights_fn())
+            self.policy_weights.apply(lambda x: x.data).update_(self.get_weights_fn())
 
     def __iter__(self) -> Iterator[TensorDictBase]:
         return self.iterator()