diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 841dcd7c4d..d6f6ae6f0f 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -709,7 +709,7 @@ def LlamaModel_fast_forward(
     if attention_mask is None:
         padding_mask = None
     # elif self.training:
-    elif attention_mask is not None and self.training:
+    elif attention_mask is not None:
         attention_mask = None
         padding_mask = None
     else: