From 72a7919e3711b7ce6eee275c50b94e2f7c4785d6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 2 Mar 2025 23:08:44 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index fcba2eb6d4..1c9998e1c9 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -279,10 +279,11 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): n_rows : int vocab_size : int n_rows, vocab_size = logits.shape + device = logits.device div, mod = divmod(vocab_size, MAX_FUSED_SIZE) n_chunks : int = div + (mod != 0) - losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + losses = torch.empty(n_rows, dtype = torch.float32, device = device) DO_SOFTCAPPING : bool = bool(logit_softcapping != 0) DO_LOGIT_SCALING : bool = bool(logit_scaling != 0) @@ -292,7 +293,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral BLOCK_SIZE, num_warps = calculate_settings(vocab_size) - logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + logsumexp = torch.empty(n_rows, dtype = torch.float32, device = device) _cross_entropy_forward[(n_rows,)]( logits, logits.stride(0), @@ -309,7 +310,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): ) else: # For large vocabs > 65336 like Gemma 256K - logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = "cuda:0") + logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = device) _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( logits, logits.stride(0),