amend

pytorch · Jul 29, 2024 · 57e4051 · 57e4051
1 parent d13e97e
commit 57e4051
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -1122,7 +1122,9 @@ def _maybe_set_truncated(self, final_rollout):
             truncated[last_step] = True
             final_rollout["next", truncated_key] = truncated
             done = final_rollout["next", _replace_last(truncated_key, "done")]
-            final_rollout["next", _replace_last(truncated_key, "done")] = done | truncated
+            final_rollout["next", _replace_last(truncated_key, "done")] = (
+                done | truncated
+            )
         return final_rollout
 
     @torch.no_grad()