Merge pull request #319 from RahulSChand/warning

Give better error message in Tinystories data loader
2026-02-06 11:26:53 +00:00 · 2023-08-19 10:00:27 -07:00 · 2023-08-19 10:00:27 -07:00 · d2a546c577
commit d2a546c577
parent bd182289c5 fbefeec1b1
1 changed files with 1 additions and 0 deletions
--- a/tinystories.py
+++ b/tinystories.py
@ -196,6 +196,7 @@ class PretokDataset(torch.utils.data.IterableDataset):
            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
        # train/test split. let's use only shard 0 for test split, rest train
        shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
+        assert len(shard_filenames)>0, f"No bin files found in {bin_dir}"
        while True:
            rng.shuffle(shard_filenames)
            for shard in shard_filenames: