diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index d33cb78409..70b0f116d5 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -279,9 +279,6 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
         n_chunks = div + (mod != 0)
         losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda: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
             BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
@@ -294,9 +291,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
                 labels,
                 VOCAB_SIZE       = vocab_size,
                 BLOCK_SIZE       = BLOCK_SIZE,
-                DO_SOFTCAPPING   = DO_SOFTCAPPING,
+                DO_SOFTCAPPING   = bool(logit_softcapping != 0),
                 SOFTCAP          = logit_softcapping,
-                DO_LOGIT_SCALING = DO_LOGIT_SCALING,
+                DO_LOGIT_SCALING = bool(logit_scaling != 0),
                 LOGIT_SCALE      = logit_scaling,
                 num_warps        = num_warps,
             )
@@ -312,9 +309,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
                 VOCAB_SIZE       = vocab_size,
                 N_CHUNKS         = n_chunks,
                 BLOCK_SIZE       = MAX_FUSED_SIZE,
-                DO_SOFTCAPPING   = DO_SOFTCAPPING,
+                DO_SOFTCAPPING   = bool(logit_softcapping != 0),
                 SOFTCAP          = logit_softcapping,
-                DO_LOGIT_SCALING = DO_LOGIT_SCALING,
+                DO_LOGIT_SCALING = bool(logit_scaling != 0),
                 LOGIT_SCALE      = logit_scaling,
                 num_warps        = 32,
             )
@@ -326,9 +323,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
         pass
 
         ctx.save_for_backward(logits, logsumexp, labels)
-        ctx.DO_SOFTCAPPING    = DO_SOFTCAPPING
+        ctx.DO_SOFTCAPPING    = bool(logit_softcapping != 0)
         ctx.logit_softcapping = logit_softcapping
-        ctx.DO_LOGIT_SCALING  = DO_LOGIT_SCALING
+        ctx.DO_LOGIT_SCALING  = bool(logit_scaling != 0)
         ctx.logit_scaling     = logit_scaling
         return losses
     pass