From 37b160d724752d761c5689228743bb1501490653 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 3 Nov 2024 15:15:58 -0800
Subject: [PATCH] CE Loss

---
 unsloth/kernels/__init__.py           | 1 +
 unsloth/kernels/cross_entropy_loss.py | 6 ++++--
 unsloth/models/_utils.py              | 1 +
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py
index 78e70a65b5..3b31f49999 100644
--- a/unsloth/kernels/__init__.py
+++ b/unsloth/kernels/__init__.py
@@ -15,6 +15,7 @@
 from .cross_entropy_loss import (
     fast_cross_entropy_loss,
     post_patch_loss_function,
+    patch_loss_functions,
 )
 from .rms_layernorm import (
     fast_rms_layernorm,
diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 92e64bfa9b..41bce690e1 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -20,7 +20,7 @@ from transformers.models.llama.modeling_llama import logger
 from packaging.version import Version
 
 from unsloth_zoo.loss_utils import (
-    patch_loss_functions,
+    patch_loss_functions as _patch_loss_functions,
     post_patch_loss_function,
 )
 
@@ -393,4 +393,6 @@ if (Version(torch.__version__) < Version("2.4.0")) and \
 pass
 
 # Patch CE Losses in transformers
-patch_loss_functions(fast_cross_entropy_loss)
+def patch_loss_functions():
+    _patch_loss_functions(fast_cross_entropy_loss)
+pass
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index f49c9db1ff..45b9569c0d 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -48,6 +48,7 @@ __all__ = [
     "patch_regional_compilation",
     "patch_layernorm",
     "patch_torch_compile",
+    "patch_loss_functions",
 ]
 
 import torch