diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 28e40487c1..266ece4ca2 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -232,7 +232,7 @@ def _cross_entropy_backward( pass # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - partial = 0.0 + partial = x if DO_SOFTCAPPING: # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x) partial = triton_tanh(x / SOFTCAP)