Update llama.py

Author: danielhanchen

unsloth/models/llama.py

@@ -858,11 +858,6 @@ def _CausalLM_fast_forward(
             logits = self.lm_head(hidden_states.to(lm_head.dtype))
         pass
         logits = logits.to(self.config.torch_dtype)
-        if self.config.final_logit_softcapping is not None:
-            logits = logits / self.config.final_logit_softcapping
-            logits = torch.tanh(logits)
-            logits = logits * self.config.final_logit_softcapping
-        pass
 
         loss = None
         if labels is not None:
@@ -876,6 +871,7 @@ def _CausalLM_fast_forward(
             loss = fast_cross_entropy_loss(
                 logits = shift_logits,
                 labels = shift_labels,
+                logit_softcapping = getattr(self.config, "final_logit_softcapping", 0),
             )
         pass