amend

pytorch · Jul 3, 2024 · 48d8855 · 48d8855
1 parent 761f28d
commit 48d8855
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torchrl/record/loggers/csv.py b/torchrl/record/loggers/csv.py
@@ -98,9 +98,8 @@ def add_video(self, tag, vid_tensor, global_step: Optional[int] = None, **kwargs
                 vid_tensor = vid_tensor.flatten(0, vid_tensor.ndim - 4)
             vid_tensor = vid_tensor.permute((0, 2, 3, 1))
             vid_tensor = vid_tensor.expand(*vid_tensor.shape[:-1], 3)
-            torchvision.io.write_video(
-                filepath, vid_tensor, fps=self.video_fps, **kwargs
-            )
+            kwargs.setdefault("fps", self.video_fps)
+            torchvision.io.write_video(filepath, vid_tensor, **kwargs)
         else:
             raise ValueError(
                 f"Unknown video format {self.video_format}. Must be one of 'pt', 'memmap' or 'mp4'."