Fix _query_videos return shapes

2024-10-23 18:18:28 +02:00 · 2024-10-23 18:18:28 +02:00 · 07570f867f
parent b8bdbc1c5b
commit 07570f867f
1 changed files with 1 additions and 1 deletions
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@ -487,7 +487,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
            frames = decode_video_frames_torchvision(
                video_path, query_ts, self.tolerance_s, self.video_backend
            )
-            item[vid_key] = frames
+            item[vid_key] = frames.squeeze(0)

        return item