fix bugs when len of traj is 1

2025-03-08 01:05:42 +08:00 · 2025-03-08 01:05:42 +08:00 · 82a2dc2415
parent 074f0ac8fe
commit 82a2dc2415
1 changed files with 1 additions and 0 deletions
--- a/lerobot/common/datasets/v21/convert_stats.py
+++ b/lerobot/common/datasets/v21/convert_stats.py
@ -40,6 +40,7 @@ def convert_episode_stats(dataset: LeRobotDataset, ep_idx: int):
        if ft["dtype"] == "video":
            # We sample only for videos
            ep_ft_data = sample_episode_video_frames(dataset, ep_idx, key)
+            ep_ft_data = ep_ft_data[None, ...] if ep_ft_data.ndim == 3 else ep_ft_data
        else:
            ep_ft_data = np.array(ep_data[key])