From 2ea161c913dbde4d0db7ffae53206961db3376c8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 14:15:08 -0800 Subject: [PATCH] CE Loss --- unsloth/kernels/__init__.py | 3 +-- unsloth/kernels/cross_entropy_loss.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) 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)