From 40fc7ec6ce1b1fde9deea0fda65f7e87ca70be41 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:23:28 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index d396538e69..09eb0854e8 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -190,7 +190,7 @@ def _cross_entropy_backward( logits_ptr , logits_row_stride : tl.constexpr(tl.int64), dloss_ptr , - dloss_row_stride , + dloss_row_stride : tl.constexpr(tl.int32), logsumexp_ptr , labels_ptr , VOCAB_SIZE : tl.constexpr(tl.int32),