From 07570f867f03b2dfef58b46e9b63153507769f96 Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Wed, 23 Oct 2024 18:18:28 +0200 Subject: [PATCH] Fix _query_videos return shapes --- lerobot/common/datasets/lerobot_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py index 0c62756e..b5e18964 100644 --- 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