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

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 811e6ccd18..1eae97ff1c 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -449,7 +449,6 @@ def LlamaAttention_fast_forward(
     else:
         # Grouped query attention
         if SDPA_HAS_GQA:
-            print("##")
             # Needs (batch_size, n_heads, seq_len, head_dim)
             # is_casual and attention_mask must not be both set!
             A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False, enable_gqa = n_groups != 1)
@@ -708,8 +707,8 @@ def LlamaModel_fast_forward(
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
-    # elif self.training:
-    elif attention_mask is None:
+    elif self.training:
+    # elif attention_mask is None:
         attention_mask = None
         padding_mask = None
     else: