Update cross_entropy_loss.py

2024-02-26 03:14:04 +11:00 · 2024-02-26 03:14:04 +11:00 · 096544eff4
commit 096544eff4
parent b5eff3bfbc
1 changed files with 1 additions and 1 deletions
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@ -247,7 +247,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
            dlosses, dlosses.stride(0),
            logsumexp,
            labels,
-            VOCAB_SIZE = n_cols,
+            VOCAB_SIZE = vocab_size,
            BLOCK_SIZE = BLOCK_SIZE,
            num_warps  = 8,
        )