Update cross_entropy_loss.py

This commit is contained in:
Daniel Han-Chen 2024-02-26 03:29:19 +11:00
commit c1dbf6710e

View file

@ -217,7 +217,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
n_rows, n_cols = logits.shape
grid = lambda meta: (n_rows, triton.cdiv(n_cols, meta["BLOCK_SIZE"]))
print(logits.stride(), dlosses.stride(), dlosses.shape)
print(logits.stride(), dlosses.stride(), dlosses.shape, dlosses)
_cross_entropy_backward[grid](
logits, logits.stride(0),
dlosses, dlosses.stride(0),