diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 86e7bd168e..3775d78e5e 100644
--- 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,
         )