Update cross_entropy_loss.py

2024-02-26 16:19:39 +11:00 · 2024-02-26 16:19:39 +11:00 · 7cf69dd9f2
commit 7cf69dd9f2
parent 1f450cd6ee
1 changed files with 1 additions and 1 deletions
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@ -251,7 +251,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
            dlosses, dlosses.stride(0),
            logsumexp,
            labels,
-            vocab_size,
+            VOCAB_SIZE = vocab_size,
            BLOCK_SIZE = BLOCK_SIZE,
            num_warps  = 8,
        )