From 0309a9fcbc47d6aeda84289762233d8851b21f4a Mon Sep 17 00:00:00 2001 From: Remi Cadene Date: Tue, 6 May 2025 15:13:50 +0000 Subject: [PATCH] Speedup data loading --- lerobot/common/datasets/lerobot_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py index 0d2800632..4df310284 100644 --- a/lerobot/common/datasets/lerobot_dataset.py +++ b/lerobot/common/datasets/lerobot_dataset.py @@ -736,7 +736,7 @@ class LeRobotDataset(torch.utils.data.Dataset): query_timestamps = {} for key in self.meta.video_keys: if query_indices is not None and key in query_indices: - timestamps = self.hf_dataset.select(query_indices[key])["timestamp"] + timestamps = self.hf_dataset[query_indices[key]]["timestamp"] query_timestamps[key] = torch.stack(timestamps).tolist() else: query_timestamps[key] = [current_ts] @@ -745,7 +745,7 @@ class LeRobotDataset(torch.utils.data.Dataset): def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict: return { - key: torch.stack(self.hf_dataset.select(q_idx)[key]) + key: torch.stack(self.hf_dataset[q_idx][key]) for key, q_idx in query_indices.items() if key not in self.meta.video_keys }