admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-16 00:59:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Removed .item from save_dataset_to_safetensors

Browse Source
This commit is contained in:
Michel Aractingi
2025-08-28 09:52:37 +02:00
parent 64a9dd3763
commit 2b03dec01f
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/artifacts/datasets/save_dataset_to_safetensors.py
+4 -4
View File
@@ -47,15 +47,15 @@ def save_dataset_to_safetensors(output_dir, repo_id="lerobot/pusht"):
)
# save 2 first frames of first episode
i = dataset.meta.episodes["dataset_from_index"][0].item()
i = dataset.meta.episodes["dataset_from_index"][0]
save_file(dataset[i], repo_dir / f"frame_{i}.safetensors")
save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")
# save 2 frames at the middle of first episode
i = int(
(
dataset.meta.episodes["dataset_to_index"][0].item()
- dataset.meta.episodes["dataset_from_index"][0].item()
dataset.meta.episodes["dataset_to_index"][0]
- dataset.meta.episodes["dataset_from_index"][0]
)
/ 2
)
@@ -63,7 +63,7 @@ def save_dataset_to_safetensors(output_dir, repo_id="lerobot/pusht"):
save_file(dataset[i + 1], repo_dir / f"frame_{i + 1}.safetensors")
# save 2 last frames of first episode
i = dataset.meta.episodes["dataset_to_index"][0].item()
i = dataset.meta.episodes["dataset_to_index"][0]
save_file(dataset[i - 2], repo_dir / f"frame_{i - 2}.safetensors")
save_file(dataset[i - 1], repo_dir / f"frame_{i - 1}.safetensors")