Update cross_entropy_loss.py

2024-10-30 16:55:24 -07:00 · 2024-10-30 16:55:24 -07:00 · 30cdf652d3
commit 30cdf652d3
parent 9f926ced28
1 changed files with 1 additions and 1 deletions
--- 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)