From cb5c182bda280515ef3db9efa0ee4b94b07008f6 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 26 Dec 2024 19:30:49 -0800
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 14da9fc426..18918a3c76 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1011,7 +1011,8 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches):
     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])
-        except (TypeError, AttributeError):
+        except Exception as exception:
+            logger.warning_once(exception)
             pass
 
     if self.args.average_tokens_across_devices: