diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 5c34df71b9..454dc27f22 100644
--- 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,
         )