From 388d4518ba20c334b002442527126f738a9629c3 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Sat, 7 Mar 2026 17:57:56 +0100 Subject: [PATCH] fix(datasets) --- src/lerobot/datasets/lerobot_dataset.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 617ac297f..f1d3487af 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -486,8 +486,8 @@ class LeRobotDataset(torch.utils.data.Dataset): self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes) # Check timestamps - timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy() - episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy() + timestamps = torch.tensor(self.hf_dataset["timestamp"]).numpy() + episode_indices = torch.tensor(self.hf_dataset["episode_index"]).numpy() ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()} check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s) @@ -667,7 +667,7 @@ class LeRobotDataset(torch.utils.data.Dataset): 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"] - query_timestamps[key] = torch.stack(timestamps).tolist() + query_timestamps[key] = torch.tensor(timestamps).tolist() else: query_timestamps[key] = [current_ts] @@ -675,7 +675,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.tensor(self.hf_dataset.select(q_idx)[key]) for key, q_idx in query_indices.items() if key not in self.meta.video_keys }