From f5f3d6794ecf70e8ece5443b6c986168d8d0c418 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 02:08:26 +1100 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 4b4bda0b4e..184b08c674 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -106,7 +106,7 @@ def _large_cross_entropy_forward( if (label_idx >= (col_idx+0)*BLOCK_SIZE) and \ (label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)): - loss = tl.load(logits_ptr + label_idx).to(tl.float32) + logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) lse = 0.0 loss = lse - logits_label # We add the final logsumexp after a reduction pass