From 6184f770e7c72303c1cd6b6de78c3aaf26d1d9bc Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 02:10:38 +1100 Subject: [PATCH] Update cross_entropy_loss.py --- 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 184b08c674..8c31abea30 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -188,8 +188,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): ) else: # For small vocabs > 65336 like Gemma - losses = torch.empty((n_rows, n_splits), dtype = torch.float32, device = "cuda") - logsumexp = torch.empty((n_rows, n_splits), dtype = torch.float32, device = "cuda") + losses = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda") + logsumexp = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda") _large_cross_entropy_forward[(n_rows, n_splits,)]( logits, logits.stride(0),