Update cross_entropy_loss.py
This commit is contained in:
parent
bc3ff0f3a6
commit
b55d7ad4b2
1 changed files with 1 additions and 1 deletions
|
|
@ -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())
|
||||
print(logits.stride(), dlosses.stride(), dlosses.shape)
|
||||
_cross_entropy_backward[grid](
|
||||
logits, logits.stride(0),
|
||||
dlosses, dlosses.stride(0),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue