Fix episode chunk

2024-10-23 00:27:14 +02:00
parent c3c0141738
commit 9dca233d7e
1 changed files with 1 additions and 4 deletions
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@@ -305,10 +305,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
        return Path(fpath)

    def get_episode_chunk(self, ep_index: int) -> int:
-        ep_chunk = ep_index // self.chunks_size
-        if ep_index > 0 and ep_index % self.chunks_size == 0:
-            ep_chunk -= 1
-        return ep_chunk
+        return ep_index // self.chunks_size

    @property
    def data_path(self) -> str: