rebased from user/rcadene/2024_04_18_episode_data_index

This commit is contained in:
Simon Alibert 2024-04-20 14:50:34 +02:00
parent 260dc3ff35
commit 7e3f284e67
1 changed files with 1 additions and 1 deletions

View File

@ -58,7 +58,7 @@ def hf_transform_to_torch(items_dict):
def load_hf_dataset(dataset_id, version, root, split) -> datasets.Dataset: def load_hf_dataset(dataset_id, version, root, split) -> datasets.Dataset:
"""hf_dataset contains all the observations, states, actions, rewards, etc.""" """hf_dataset contains all the observations, states, actions, rewards, etc."""
if root is not None: if root is not None:
hf_dataset = load_from_disk(Path(root) / dataset_id / split) hf_dataset = load_from_disk(str(Path(root) / dataset_id / split))
else: else:
# TODO(rcadene): remove dataset_id everywhere and use repo_id instead # TODO(rcadene): remove dataset_id everywhere and use repo_id instead
repo_id = f"lerobot/{dataset_id}" repo_id = f"lerobot/{dataset_id}"