From 74a349e12faa06c90dd0cce85feb8509e4101357 Mon Sep 17 00:00:00 2001 From: Nole Lin Date: Sat, 14 Feb 2026 23:46:32 -0800 Subject: [PATCH 1/2] bug fix for multi-GPU --- unsloth/kernels/cross_entropy_loss.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index cdd5b1cf12..65c502eee0 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -438,6 +438,7 @@ def fast_cross_entropy_loss( batch, seq_len, d = logits.shape assert labels.shape == (batch, seq_len) + device = logits.device loss = Fast_CrossEntropyLoss.apply( logits.view(batch * seq_len, d), labels.view(-1), @@ -446,6 +447,8 @@ def fast_cross_entropy_loss( ) if n_items is None: n_items = torch.count_nonzero(labels != -100) + loss = loss.to(device) + n_items = n_items.to(device) return loss.sum() / n_items From 5f687612f82f519451a61d8800e4dd87d9b0d719 Mon Sep 17 00:00:00 2001 From: nole69 <89107458+nole69@users.noreply.github.com> Date: Sun, 15 Feb 2026 00:03:28 -0800 Subject: [PATCH 2/2] Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- unsloth/kernels/cross_entropy_loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 65c502eee0..d92229314f 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -447,8 +447,8 @@ def fast_cross_entropy_loss( ) if n_items is None: n_items = torch.count_nonzero(labels != -100) - loss = loss.to(device) - n_items = n_items.to(device) + if torch.is_tensor(n_items): + n_items = n_items.to(device) return loss.sum() / n_items