diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 2010a4cbe4..5c34df71b9 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -218,6 +218,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
                 logsumexp,
                 labels,
                 VOCAB_SIZE = vocab_size,
+                N_CHUNKS   = n_chunks,
                 BLOCK_SIZE = MAX_FUSED_SIZE,
                 num_warps  = 32,
             )