diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 5236e51985..f5a015073c 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -20,9 +20,8 @@ from transformers.models.llama.modeling_llama import logger
 from packaging.version import Version
 
 from unsloth_zoo.loss_utils import (
-    causal_loss_function,
-    transformers_losses_patcher,
-    patch_loss_function,
+    patch_loss_functions,
+    post_patch_loss_function,
 )