fixing datasets without max sequences

stanford-crfm · Feb 15, 2025 · 2eaec4f · 2eaec4f
1 parent 5b7ddbb
commit 2eaec4f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/levanter/data/text.py b/src/levanter/data/text.py
@@ -1284,6 +1284,8 @@ def shuffle_ds(ds, key):
         for name, ds in token_datasets.items():
             if self.max_sequences_dict is not None and name in self.max_sequences_dict:
                 train_token_datasets[name] = ds.slice_dataset(end_index=self.max_sequences_dict[name])
+            else:
+                train_token_datasets[name] = ds
 
         self.validation_token_datasets = {}
         for name, ds in token_datasets.items():