From a838f6ece6c9b8035b42f335283840dbb2bd2b15 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Dec 2024 19:12:13 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 512812cb7c..14da9fc426 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1008,8 +1008,6 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches): batch_samples += [next(epoch_iterator)] except StopIteration: break - - print("NUM_ITMES = ", num_items_in_batch, type(batch_samples)) if len(batch_samples) > 0 and "labels" in batch_samples[0]: try: num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples]) @@ -1022,6 +1020,8 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches): if torch.is_tensor(num_items_in_batch): num_items_in_batch = num_items_in_batch.item() + print("NUM_ITMES = ", num_items_in_batch, type(batch_samples)) + return batch_samples, num_items_in_batch # def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):