diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 088450b9e7..0b567b023a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -707,8 +707,8 @@ def LlamaModel_fast_forward(
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
-    elif attention_mask is None and self.training:
-        # attention_mask = None
+    elif self.training:
+        attention_mask = None
         padding_mask = None
     else:
         # if 0 in attention_mask: