From 336f3b9c24bafd376ca30dde4cd2f06f301b377a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Sep 2024 00:01:56 -0700
Subject: [PATCH] Update cross_entropy_loss.py

---
 unsloth/kernels/cross_entropy_loss.py | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index e5ba8b55f1..a09f167db7 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -377,17 +377,17 @@ def fast_cross_entropy_loss(
 pass
 
 
-from transformers.models.llama.modeling_llama import (
-    LlamaForCausalLM,
-    CausalLMOutputWithPast,
-    Optional,
-    Union,
-    Cache,
-    List,
-    Tuple,
-)
+from transformers.models.llama.modeling_llama import LlamaForCausalLM
 def patch_llama_for_causal_lm():
     import transformers.models.llama.modeling_llama
+    from transformers.models.llama.modeling_llama import (
+        CausalLMOutputWithPast,
+        Optional,
+        Union,
+        Cache,
+        List,
+        Tuple,
+    )
     import inspect, re
     function = inspect.getsource(transformers.models.llama.modeling_llama.LlamaForCausalLM.forward)
     function = function.split("\n")
@@ -448,9 +448,8 @@ def patch_llama_for_causal_lm():
 
     patched_function = f"class Unsloth_LlamaForCausalLM(LlamaForCausalLM):\n"\
     f"    {function}\n"
-
-    print(patched_function)
-    exec(patched_function)
+    
+    exec(patched_function, globals())
     transformers.models.llama.modeling_llama.LlamaForCausalLM = Unsloth_LlamaForCausalLM
     return
 pass