Update

[ghstack-poisoned]
pytorch · Feb 24, 2025 · 8ae33aa · 8ae33aa
1 parent 348558c
commit 8ae33aa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -846,7 +846,8 @@ def _maybe_make_final_rollout(self, make_rollout: bool):
                     if key in self._final_rollout.keys(True):
                         continue
                     self._final_rollout.set(key, spec.zero())
-
+        elif not make_rollout and hasattr(self.policy, "out_keys") and self.policy.out_keys:
+            self._policy_output_keys = list(self.policy.out_keys)
         else:
             if make_rollout:
                 # otherwise, we perform a small number of steps with the policy to