From 8c88b29d74e792ffc115ebc232c29b73325f18d8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 13 Feb 2025 17:22:39 -0800
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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: