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

---
 unsloth/kernels/cross_entropy_loss.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index e791664a6f..e5ba8b55f1 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
+from transformers.models.llama.modeling_llama import (
+    LlamaForCausalLM,
+    CausalLMOutputWithPast,
+    Optional,
+    Union,
+    Cache,
+    List,
+    Tuple,
+)
 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")