fix the fixes 3

vmoens · vmoens · commit 3428d3f0d578 · 2025-11-10T17:49:39.000Z
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -2783,8 +2783,9 @@ def _setup_multi_policy_and_weights(
                                 )
 
                 # Extract weights from policy
+                # Use .data to avoid gradient tracking (can't serialize tensors with requires_grad)
                 weights = (
-                    TensorDict.from_module(policy)
+                    TensorDict.from_module(policy, as_module=True).data
                     if isinstance(policy, nn.Module)
                     else TensorDict()
                 )

Original file line number	Diff line number	Diff line change
`@@ -2783,8 +2783,9 @@ def _setup_multi_policy_and_weights(`
`2783`	`2783`	`)`
`2784`	`2784`
`2785`	`2785`	`# Extract weights from policy`
	`2786`	`+ # Use .data to avoid gradient tracking (can't serialize tensors with requires_grad)`
`2786`	`2787`	`weights = (`
`2787`		`- TensorDict.from_module(policy)`
	`2788`	`+ TensorDict.from_module(policy, as_module=True).data`
`2788`	`2789`	`if isinstance(policy, nn.Module)`
`2789`	`2790`	`else TensorDict()`
`2790`	`2791`	`)`