From e3c9482cfb6d8f06faa092cebc38344638895a30 Mon Sep 17 00:00:00 2001
From: nole69 <89107458+nole69@users.noreply.github.com>
Date: Sun, 15 Feb 2026 01:09:40 -0800
Subject: [PATCH] [FIX] Move loss and n_items to logits device in
 fast_cross_entropy_loss loss for multi-GPU support (#4063)

* bug fix for multi-GPU

* Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 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..d92229314f 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)
+    if torch.is_tensor(n_items):
+        n_items = n_items.to(device)
     return loss.sum() / n_items