diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 29de5858de..32b1daaa01 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1009,42 +1009,34 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
     f = model.base_model.model.forward if hasattr(model, "base_model") else model.forward
     has_kwargs = tuple(inspect.signature(f).parameters.values())[-1].kind == inspect._VAR_KEYWORD
 
+    # Iterate to find all batches
     for _ in range(num_batches):
         try:
             batch_samples += [next(epoch_iterator)]
         except StopIteration:
             break
+    pass
+
+    # Get num_items_in_batch
     if has_kwargs and 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])
+            num_items_in_batch = sum(
+                [(x["labels"][..., 1:] != -100).sum() for x in batch_samples]
+            )
+            # num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
 
             if self.args.average_tokens_across_devices:
                 num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item()
 
             if torch.is_tensor(num_items_in_batch):
                 num_items_in_batch = num_items_in_batch.item()
+
         except Exception as exception:
             logger.warning_once(exception)
     pass
-    return batch_samples, num_items_in_batch
 
-# def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
-#     batch_samples = []
-#     num_items_in_batch = None
-#     for _ in range(num_batches):
-#         try:
-#             batch_samples += [next(epoch_iterator)]
-#         except StopIteration:
-#             break
-#     if len(batch_samples) > 0 and "labels" in batch_samples[0]:
-#         try:
-#             num_items_in_batch = sum(
-#                 [torch.count_nonzero(x["labels"][..., 1:] != -100) for x in batch_samples]
-#             )
-#         except TypeError:
-#             pass
-#     return batch_samples, num_items_in_batch
-# pass
+    return batch_samples, num_items_in_batch
+pass
 
 
 def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):