diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index e04c573c6c..fbc6d53afb 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -706,9 +706,10 @@ def LlamaModel_fast_forward(
     pass
 
     # Ignore attention_mask
+    print(attention_mask, attention_mask.shape, attention_mask.dtype)
     if attention_mask is None:
         padding_mask = None
-    elif self.training:
+    elif attention_mask is None and self.training:
         # attention_mask = None
         padding_mask = None
     else: