From ca37a2e9501f8d39e85752e4f940aeb3007d7c25 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 3 Nov 2024 21:27:37 -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 d396538e69..f92ddbd05d 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -32,7 +32,7 @@ from unsloth_zoo.loss_utils import (
 @triton.jit
 def _cross_entropy_forward(
     logits_ptr        ,
-    logits_row_stride : tl.constexpr(tl.int64),
+    logits_row_stride ,
     loss_ptr          ,
     logsumexp_ptr     ,
     labels_ptr        ,
@@ -105,7 +105,7 @@ pass
 @triton.jit
 def _chunked_cross_entropy_forward(
     logits_ptr        ,
-    logits_row_stride : tl.constexpr(tl.int64),
+    logits_row_stride ,
     loss_ptr          ,
     logsumexp_ptr     ,
     labels_ptr        ,
@@ -188,7 +188,7 @@ pass
 @triton.jit
 def _cross_entropy_backward(
     logits_ptr        ,
-    logits_row_stride : tl.constexpr(tl.int64),
+    logits_row_stride ,
     dloss_ptr         ,
     dloss_row_stride  ,
     logsumexp_ptr     ,