From 8034edf925203a329424595c58ee81eefd606e69 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 13 Feb 2025 14:42:38 -0800
Subject: [PATCH] Update llama.py

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index d6916814ae..ec6706e515 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -707,7 +707,8 @@ def LlamaModel_fast_forward(
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
-    elif attention_mask is not None and self.training:
+    elif self.training:
+    # elif attention_mask is not None and self.training:
         attention_mask = None
         padding_mask = None
     else: