From cf2a714d747dd6a751e9cb49d8f2aa30d47232e3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 5 Nov 2024 14:43:49 -0800
Subject: [PATCH] Update cross_entropy_loss.py

---
 unsloth/kernels/cross_entropy_loss.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 13d90baafe..e5136f8412 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -181,10 +181,10 @@ def _chunked_cross_entropy_forward(
 pass
 
 
-# @triton.heuristics({
-#     "DO_SOFTCAPPING":   lambda args: bool(args["DO_SOFTCAPPING"  ]),
-#     "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
-# })
+@triton.heuristics({
+    "DO_SOFTCAPPING":   lambda args: bool(args["DO_SOFTCAPPING"  ]),
+    "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
+})
 @triton.jit
 def _cross_entropy_backward(
     logits_ptr        ,
@@ -337,6 +337,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
         return losses
     pass
 
+
+    @torch.compiler.disable
     @staticmethod
     def backward(ctx, dlosses):
         logits, logsumexp, labels = ctx.saved_tensors