diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index abe08d7d06..f0193c74d8 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -36,7 +36,7 @@ def _cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , SOFTCAP , @@ -109,7 +109,7 @@ def _chunked_cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , N_CHUNKS , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , @@ -193,7 +193,7 @@ def _cross_entropy_backward( dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , SOFTCAP ,