diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index f2a86fde28..b96a9c4f4b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -486,6 +486,15 @@ def LlamaModel_fast_forward(
     if inputs_embeds is None:
         inputs_embeds = self.embed_tokens(input_ids)
 
+    # Fix up attention mask by setting elements to 0
+    # Specifically for DPO
+    if self._has_no_labels and attention_mask is not None:
+        inputs_requires_grad = inputs_embeds.requires_grad
+        if inputs_requires_grad: inputs_embeds.requires_grad_(False)
+        inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
+        if inputs_requires_grad: inputs_embeds.requires_grad_(True)
+    pass
+
     # Ignore attention_mask
     if attention_mask is None:
         padding_mask = None
@@ -518,17 +527,6 @@ def LlamaModel_fast_forward(
             use_cache = False
     pass
 
-    # Fix up attention mask by setting elements to 0
-    # Specifically for DPO
-    print(self._has_no_labels, attention_mask)
-    if self._has_no_labels and attention_mask is not None:
-        inputs_requires_grad = hidden_states.requires_grad
-        if inputs_requires_grad: hidden_states.requires_grad_(False)
-        hidden_states *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
-        print(1)
-        if inputs_requires_grad: hidden_states.requires_grad_(True)
-    pass
-
     # decoder layers
     all_hidden_states = () if output_hidden_states else None
     all_self_attns = () if output_attentions else None