Update cross_entropy_loss.py

2024-02-26 03:12:23 +11:00 · 2024-02-26 03:12:23 +11:00 · b5eff3bfbc
commit b5eff3bfbc
parent fdaa9493be
1 changed files with 1 additions and 0 deletions
--- 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,
            )