diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 9725f624ad..2a7532a991 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1026,8 +1026,6 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
         except Exception as exception:
             logger.warning_once(exception)
     pass
-
-    print(batch_samples, num_items_in_batch)
     return batch_samples, num_items_in_batch
 
 # def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
@@ -1051,6 +1049,9 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
 
 def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):
     if "num_items_in_batch" in kwargs:
+        if kwargs["num_items_in_batch"] is None:
+            # Remove it since the model does not support it!
+            kwargs.pop("num_items_in_batch", None)
         if "num_items_in_batch" not in inputs:
             inputs["num_items_in_batch"] = kwargs["num_items_in_batch"]
         pass