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