diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index af144f01a1..0b567b023a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -707,9 +707,9 @@ def LlamaModel_fast_forward(
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
-    # elif self.training:
-    #     attention_mask = None
-    #     padding_mask = None
+    elif self.training:
+        attention_mask = None
+        padding_mask = None
     else:
         # if 0 in attention_mask:
         #     padding_mask = attention_mask
@@ -1161,7 +1161,7 @@ def CausalLM_fast_forward(fast_forward_inference):
         if not return_dict:
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
-        print("***", logits.dtype, logits.shape)
+
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,