From 01a1060e1135627c0269b6c24840b691f9cb32f2 Mon Sep 17 00:00:00 2001 From: ChanderG Date: Tue, 18 Jun 2024 18:29:03 +0530 Subject: [PATCH] bugfix: make mock data loader usable Signed-off-by: ChanderG --- src/instructlab/training/token_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/instructlab/training/token_dataset.py b/src/instructlab/training/token_dataset.py index c515888d..374714ae 100644 --- 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):