From c7f02840d573441adcaae26e27f1bd1ea9f3cc0a Mon Sep 17 00:00:00 2001
From: pluesclues <136766175+pluesclues@users.noreply.github.com>
Date: Tue, 7 Oct 2025 23:55:04 -0400
Subject: [PATCH 1/4] Update compiler.py, temporary patch for the logits and
 stuff

---
 unsloth_zoo/compiler.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index 3125db73b..3a5a36aaf 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1339,6 +1339,15 @@ def apply_fused_lm_head(forward, module = None):
         forward = forward.replace(", **)", ")")
         forward = forward.replace(",**)", ")")
         forward = forward.replace(",** )", ")")
+
+        line_to_replace = "        logits = EMPTY_LOGITS"
+
+        replacement_block = """        if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1':
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+        else:
+            logits = EMPTY_LOGITS"""
+        forward = forward.replace(line_to_replace, replacement_block)
+    
         # print(forward)
         return forward
     pass

From 0ad7ef3f6a4f1322e41db6c1e53f989d54f9536b Mon Sep 17 00:00:00 2001
From: pluesclues <136766175+pluesclues@users.noreply.github.com>
Date: Mon, 13 Oct 2025 19:40:12 -0400
Subject: [PATCH 2/4] Update compiler.py, commit new cahnges

---
 unsloth_zoo/compiler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index 3a5a36aaf..5c4058a50 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1311,7 +1311,7 @@ def apply_fused_lm_head(forward, module = None):
         replacement = cross_entropy_replacement.strip().split("\n")
         replacement = "\n".join((len(spaces)-4)*" " + x for x in replacement)
         replacement = \
-            "logits = EMPTY_LOGITS\n" + \
+            "logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS\n" + \
             (len(spaces)-4)*" " + "loss = None\n" + \
             replacement + "\n"
         try:

From e1736602d7d28dfe3fb6c45d9146cb818f5321ea Mon Sep 17 00:00:00 2001
From: pluesclues <136766175+pluesclues@users.noreply.github.com>
Date: Mon, 13 Oct 2025 19:45:49 -0400
Subject: [PATCH 3/4] Update compiler.py. simplified code and pr

---
 unsloth_zoo/compiler.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index 5c4058a50..4a9ac62fe 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1339,15 +1339,6 @@ def apply_fused_lm_head(forward, module = None):
         forward = forward.replace(", **)", ")")
         forward = forward.replace(",**)", ")")
         forward = forward.replace(",** )", ")")
-
-        line_to_replace = "        logits = EMPTY_LOGITS"
-
-        replacement_block = """        if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1':
-            logits = self.lm_head(hidden_states[:, slice_indices, :])
-        else:
-            logits = EMPTY_LOGITS"""
-        forward = forward.replace(line_to_replace, replacement_block)
-    
         # print(forward)
         return forward
     pass

From a54208199c02208a4f3f7d3d06abf6b202e4bdaa Mon Sep 17 00:00:00 2001
From: pluesclues <136766175+pluesclues@users.noreply.github.com>
Date: Mon, 13 Oct 2025 19:48:59 -0400
Subject: [PATCH 4/4] Update compiler.py, made logit one liner in case return
 logits and neabled

---
 unsloth_zoo/compiler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth_zoo/compiler.py b/unsloth_zoo/compiler.py
index 4a9ac62fe..83c7a444c 100644
--- a/unsloth_zoo/compiler.py
+++ b/unsloth_zoo/compiler.py
@@ -1326,7 +1326,7 @@ def apply_fused_lm_head(forward, module = None):
         # Return logits back
         if "logits = outputs.logits" in cross_entropy_find:
             forward = forward.replace(
-                "logits = EMPTY_LOGITS",
+                "logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS",
                 "logits = outputs.logits",
             )
         # Fix vocab_size = (vocab_size=