From be21b4b6f6253a1e71d96ada960c97f143ecac7a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 20:03:03 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 70b0f116d5..347f9eb9da 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -346,9 +346,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): labels, VOCAB_SIZE = vocab_size, BLOCK_SIZE = BLOCK_SIZE, - DO_SOFTCAPPING = ctx.DO_SOFTCAPPING, + DO_SOFTCAPPING = bool(ctx.DO_SOFTCAPPING), SOFTCAP = ctx.logit_softcapping, - DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING, + DO_LOGIT_SCALING = bool(ctx.DO_LOGIT_SCALING), LOGIT_SCALE = ctx.logit_scaling, num_warps = 8, )