From 5eddfd8cfe409b70d28ac4a5ba48afecbd97d77c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:42:13 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index a7ef164df0..b4780db278 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(tl.int32), + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING , SOFTCAP , @@ -109,8 +109,8 @@ def _chunked_cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr(tl.int32), - N_CHUNKS : tl.constexpr(tl.int32), + VOCAB_SIZE , + N_CHUNKS , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING , SOFTCAP , @@ -193,7 +193,7 @@ def _cross_entropy_backward( dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr(tl.int32), + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING , SOFTCAP ,