admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-25 18:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(datasets)

Browse Source
This commit is contained in:
CarolinePascal
2026-03-07 17:57:56 +01:00
parent 232dbe4176
commit 388d4518ba
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/datasets/lerobot_dataset.py
+4 -4
View File
@@ -486,8 +486,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes) self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
# Check timestamps # Check timestamps
timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy() timestamps = torch.tensor(self.hf_dataset["timestamp"]).numpy()
episode_indices = torch.stack(self.hf_dataset["episode_index"]).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()} 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) 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: for key in self.meta.video_keys:
if query_indices is not None and key in query_indices: if query_indices is not None and key in query_indices:
timestamps = self.hf_dataset.select(query_indices[key])["timestamp"] timestamps = self.hf_dataset.select(query_indices[key])["timestamp"]
query_timestamps[key] = torch.stack(timestamps).tolist() query_timestamps[key] = torch.tensor(timestamps).tolist()
else: else:
query_timestamps[key] = [current_ts] 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: def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
return { 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() for key, q_idx in query_indices.items()
if key not in self.meta.video_keys if key not in self.meta.video_keys
} }