diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 69f36f0d46..0be3bfcc1d 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1128,6 +1128,7 @@ def unsloth_compile_transformers(
     debug                   = False,
     import_from_cache       = False,
     disable                 = False,
+    return_logits           = False,
 ):
     if disable: return
     model_types = get_transformers_model_type(
@@ -1158,6 +1159,7 @@ def unsloth_compile_transformers(
             debug                  = debug,
             import_from_cache      = import_from_cache,
             disable                = disable,
+            return_logits          = return_logits,
         )
     pass
     return model_types
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 0256fc1830..3ec59b339a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -974,13 +974,15 @@ def CausalLM_fast_forward(fast_forward_inference):
         logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
         logit_scaling     = getattr(self.config, "logit_scale", 0)
 
+        print(kwargs)
+
         if bsz == 1 and q_len == 1:
             logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
             logits = logits.unsqueeze(0).unsqueeze(0)
         elif num_logits_to_keep != 0:
             logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype))
         else:
-            if HAS_CUT_CROSS_ENTROPY and labels is not None:
+            if "return_logits" not in kwargs and HAS_CUT_CROSS_ENTROPY and labels is not None:
                 n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None)
                 loss = fused_linear_cross_entropy(
                     hidden_states      = hidden_states,
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 452ec5ccb3..fb97e74622 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -350,7 +350,7 @@ class FastVisionModel(FastBaseVisionModel):
         *args, **kwargs,
     ):
         if token is None: token = get_token()
-
+        
         patch_compiled_autograd()
         patch_compiling_bitsandbytes()
         if use_gradient_checkpointing == "unsloth":