Update cross_entropy_loss.py
This commit is contained in:
parent
e1b8dcf536
commit
b2cdb409bd
1 changed files with 3 additions and 3 deletions
|
|
@ -36,7 +36,7 @@ def _cross_entropy_forward(
|
|||
loss_ptr ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
VOCAB_SIZE ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
SOFTCAP ,
|
||||
|
|
@ -109,7 +109,7 @@ def _chunked_cross_entropy_forward(
|
|||
loss_ptr ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
VOCAB_SIZE ,
|
||||
N_CHUNKS ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
|
|
@ -193,7 +193,7 @@ def _cross_entropy_backward(
|
|||
dloss_row_stride ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
VOCAB_SIZE ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
SOFTCAP ,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue