From b2cdb409bdc633709b0102e131cfa132425b91ed Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 4 Nov 2024 00:32:35 -0800
Subject: [PATCH] Update cross_entropy_loss.py

---
 unsloth/kernels/cross_entropy_loss.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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           ,