From 8b3e9c2ff0b3e5c6fb11d2913023a1d9fa069324 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 Nov 2024 21:07:50 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index cc3dbb1d87..f82defd405 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -400,6 +400,6 @@ if (Version(torch.__version__) < Version("2.4.0")) and \ pass # Patch CE Losses in transformers -def patch_loss_functions(): - _patch_loss_functions(fast_cross_entropy_loss) +def patch_loss_functions(torch_compile = True): + _patch_loss_functions(fast_cross_entropy_loss, torch_compile = torch_compile) pass