From 8aefcd0b5000efcd2406e8b56bea6770d3ed9f82 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 30 Oct 2024 17:35:20 -0700 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 9f16e8e605..8cbdbf2a2c 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -273,8 +273,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): n_chunks = div + (mod != 0) losses = torch.empty(n_rows, dtype = torch.float32, device = logits.device) - DO_SOFTCAPPING = logit_softcapping != 0 - DO_LOGIT_SCALING = logit_scaling != 0 + DO_SOFTCAPPING : bool = bool(logit_softcapping != 0) + DO_LOGIT_SCALING : bool = bool(logit_scaling != 0) if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral