From 2a85cd9efc42b10b25f18a17aa7089219f69752e Mon Sep 17 00:00:00 2001 From: Matteo Bettini Date: Wed, 24 Apr 2024 11:15:43 +0100 Subject: [PATCH] [Nit] Stack video frames before transpose --- benchmarl/experiment/logger.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/benchmarl/experiment/logger.py b/benchmarl/experiment/logger.py index 92632ebf..740389fd 100644 --- 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: