diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index e41a45b266..de7f7bc49b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -897,9 +897,15 @@ def CausalLM_fast_forward(fast_forward_inference):
                 logit_softcapping = logit_softcapping,
             )
         elif logit_softcapping != 0:
-            logits *= (1.0 / logit_softcapping)
-            logits = torch.tanh(logits, out = logits if not logits.requires_grad else None)
-            logits *= logit_softcapping
+            if logits.requires_grad:
+                logits = (1.0 / logit_softcapping) * logits
+                logits = torch.tanh(logits)
+                logits = logit_softcapping * logits
+            else:
+                logits *= (1.0 / logit_softcapping)
+                torch.tanh(logits, out = logits)
+                logits *= logit_softcapping
+            pass
         pass
 
         if not return_dict: