[Nit] Stack video frames before transpose

facebookresearch · Apr 24, 2024 · 2a85cd9 · 2a85cd9
1 parent 1088fc3
commit 2a85cd9
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/benchmarl/experiment/logger.py b/benchmarl/experiment/logger.py
@@ -224,10 +224,11 @@ def log_evaluation(
 
         self.log(to_log, step=step)
         if video_frames is not None and rollouts[0].batch_size[0] > 1:
+            video_frames = np.stack(
+                video_frames[: rollouts[0].batch_size[0] - 1], axis=0
+            )
             vid = torch.tensor(
-                np.transpose(
-                    video_frames[: rollouts[0].batch_size[0] - 1], (0, 3, 1, 2)
-                ),
+                np.transpose(video_frames, (0, 3, 1, 2)),
                 dtype=torch.uint8,
             ).unsqueeze(0)
             for logger in self.loggers: