diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 7cb5b2e39e..900767aa0e 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -261,11 +261,12 @@ def grpo_trainer__get_per_token_logps(function_name, function):
         os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
         with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):
             # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
-            print("input_ids Unsloth 264", input_ids.shape)
-            print("logits_to_keep Unsloth 264", logits_to_keep)
-            hidden_states = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
-            print("hidden_states Unsloth 264", hidden_states.shape)
-            #logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
+            hidden_states = model(
+                input_ids = input_ids,
+                attention_mask = attention_mask,
+                logits_to_keep = logits_to_keep + 1,
+            ).logits
+            # logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
             return hidden_states
             # input_ids = input_ids[:, -logits_to_keep:]
             # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
@@ -318,14 +319,8 @@ def grpo_trainer_compute_loss(function_name, function):
         logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens
         _input_ids = input_ids
         _logits_to_keep = logits_to_keep
-        print("prompt_mask Unsloth 320", prompt_mask.shape)
-        print("completion_mask Unsloth 320", completion_mask.shape)
-        print("input_ids Unsloth 320", input_ids.shape)
-        print("logits_to_keep Unsloth 320", logits_to_keep)
 
         per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
-        if per_token_logps is not None:
-            print("per_token_logps Unsloth 320", per_token_logps.shape)
 
         # Compute the KL divergence between the model and the reference model
         # _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves.
@@ -333,25 +328,16 @@ def grpo_trainer_compute_loss(function_name, function):
         if self.beta != 0.0:
             with torch.inference_mode(), model.disable_adapter():
                 ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
-                if ref_per_token_logps is not None:
-                    print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
         else:
             ref_per_token_logps = None
         # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
         # x - x.detach() allows for preserving gradients from x
         advantages = inputs["advantages"]
-        print("advantages Unsloth 320", advantages.shape)
         # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
         # per_token_loss = -(per_token_loss - self.beta * per_token_kl)
         # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
         old_hidden_states = inputs.get("old_per_token_logps", None)
-        if old_hidden_states is not None:
-            print("old_hidden_states Unsloth 320", old_hidden_states.shape)
-
-        print("input_ids Unsloth 320", input_ids.shape)
-        print("logits_to_keep Unsloth 320", logits_to_keep)
         input_ids = input_ids[:, -logits_to_keep:]
-        print("input_ids Unsloth 320", input_ids.shape)
 
         # Get logit softcapping and logit scale
         logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) # Gemma
@@ -365,14 +351,9 @@ def grpo_trainer_compute_loss(function_name, function):
         if per_token_logps is not None:
 
             if ref_per_token_logps is not None:
-                print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
                 ref_per_token_logps = ref_per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
-                print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
-            
-            print("per_token_logps Unsloth 320", per_token_logps.shape)
             per_token_logps = per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
-            print("per_token_logps Unsloth 320", per_token_logps.shape)
-            
+
             loss, completion_length, mean_kl = grpo_compute_loss_slow(
                 ref_per_token_logps,
                 per_token_logps,
@@ -428,13 +409,12 @@ def grpo_trainer_compute_loss(function_name, function):
                     logit_scale_divide = logit_scale_divide,
                     attention_mask = attention_mask,
                 )
-
+            pass
+        pass
         # Log the metrics
         # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
-
         # mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
         # self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
-
         if "train" in self._metrics:
             mode = "eval" if self.control.should_evaluate else "train"
             self._metrics[mode]["completion_length"].append(completion_length.item())