Update cross_entropy_loss.py

2024-10-30 16:46:33 -07:00 · 2024-10-30 16:46:33 -07:00 · 895a5364f0
commit 895a5364f0
parent 6ef574b58a
1 changed files with 1 additions and 0 deletions
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@ -232,6 +232,7 @@ def _cross_entropy_backward(
    pass

    # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+    partial = 0.0
    if DO_SOFTCAPPING:
        # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x)
        partial = triton_tanh(x / SOFTCAP)