From cd42f4d772d11e1da374873399ebdb65fd42b6f2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 24 Sep 2024 23:55:25 -0700
Subject: [PATCH] Update cross_entropy_loss.py

---
 unsloth/kernels/cross_entropy_loss.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 58104e4dfa..d06bf69913 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -440,10 +440,12 @@ def patch_llama_for_causal_lm():
         function[ function.find("    if not return_dict"):]
     function = function.replace("logits = logits.float()", "\n")
 
+    # Missed spaces
     function = function.split("\n")
-    function = [" "*4 + x for x in function]
+    # Not the first one though!
+    function = [function[0]] + [" "*4 + x for x in function[1:]]
     function = "\n".join(function)
-    
+
     patched_function = f"class Unsloth_LlamaForCausalLM(LlamaForCausalLM):\n"\
     f"    {function}\n"