diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index bcc26b36a5..d5dd7833d5 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -489,7 +489,7 @@ def LlamaModel_fast_forward(
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
-    elif False:#self.training:
+    elif self.training:
         attention_mask = None
         padding_mask = None
     else:
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 62d628de0a..2410572174 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -90,7 +90,7 @@ def MistralAttention_fast_forward(
     past_key_value = (K, V) if use_cache else None
 
     # Attention module
-    if (attention_mask is None and not HAS_FLASH_ATTENTION):
+    if (not HAS_FLASH_ATTENTION):
         # Xformers memory efficient attention
         Q = Q.transpose(1, 2)
         K = K.transpose(1, 2)