Fix chaining in MultiLerobotDataset (#233)

2024-05-31 09:03:28 +01:00
parent 57fb5fe8a6
commit 633115d861
2 changed files with 9 additions and 1 deletions
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@@ -371,6 +371,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
            if idx >= start_idx + dataset.num_samples:
                start_idx += dataset.num_samples
                dataset_idx += 1
+                continue
            break
        else:
            raise AssertionError("We expect the loop to break out as long as the index is within bounds.")