Update

[ghstack-poisoned]

Author: vmoens

test/test_collector.py

@@ -162,7 +162,7 @@ def forward(self, observation):
         output = self.linear(observation)
         if self.multiple_outputs:
             return output, output.sum(), output.min(), output.max()
-        return self.linear(observation)
+        return output
 
 
 class UnwrappablePolicy(nn.Module):

torchrl/collectors/collectors.py

@@ -486,6 +486,9 @@ def update_policy_weights_(
 
                 strategy = WeightStrategy(extract_as="tensordict")
                 weights = strategy.extract_weights(self._original_policy)
+                # Cast weights to the policy device before applying
+                if self.policy_device is not None:
+                    weights = weights.to(self.policy_device)
                 strategy.apply_weights(self.policy, weights)
             # Otherwise, no action needed - policy is local and changes are immediately visible