ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-24 02:10:38 +11:00
commit 6184f770e7
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/cross_entropy_loss.py
View file

@ -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),