Merge pull request #44 from ChanderG/bugfix-mock-dataset

bugfix: make mock data loader usable
instructlab · Jun 18, 2024 · 1cb03c5 · 1cb03c5
2 parents 25cea8d + 01a1060
commit 1cb03c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/instructlab/training/token_dataset.py b/src/instructlab/training/token_dataset.py
@@ -63,7 +63,7 @@ def __getitem__(self, idx):
         }
 
     def get_lengths(self):
-        np.array([len(self.input_ids[0])] * len(self.input_ids))
+        return np.array([len(self.input_ids[0])] * len(self.input_ids))
 
 
 def make_collate_fn(pad_token_id, is_granite=False, max_batch_len=60000):