diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py
index df7589719d..78e70a65b5 100644
--- a/unsloth/kernels/__init__.py
+++ b/unsloth/kernels/__init__.py
@@ -14,8 +14,7 @@
 
 from .cross_entropy_loss import (
     fast_cross_entropy_loss,
-    patch_losses,
-    patch_loss_function,
+    post_patch_loss_function,
 )
 from .rms_layernorm import (
     fast_rms_layernorm,
diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index f5a015073c..92e64bfa9b 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -393,4 +393,4 @@ if (Version(torch.__version__) < Version("2.4.0")) and \
 pass
 
 # Patch CE Losses in transformers
-patch_losses = transformers_losses_patcher(causal_loss_function(fast_cross_entropy_loss))
+patch_loss_functions(fast_cross_entropy_loss)