Update

[ghstack-poisoned]
pytorch · Sep 24, 2024 · b18e666 · b18e666
1 parent 8117858
commit b18e666
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -230,7 +230,7 @@ def map_weight(
             policy = deepcopy(policy)
 
         param_and_buf.apply(
-            functools.partial(map_weight, device=policy_device),
+            functools.partial(map_weight),
             filter_empty=False,
         ).to_module(policy)
         return policy, get_original_weights