Update cross_entropy_loss.py

This commit is contained in:
Daniel Han 2024-10-30 15:46:09 -07:00
commit 8549887cd4

View file

@ -31,7 +31,7 @@ def _cross_entropy_forward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -105,7 +105,7 @@ def _chunked_cross_entropy_forward(
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -188,7 +188,7 @@ def _cross_entropy_backward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,