diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 5158019132..8436ca0dc9 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -393,12 +393,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
 
     if match:
         indent = match.group(1)
-        function = function.replace(
-            f"{indent}return output",
-            f"""{indent}if not _was_training:
-    {indent}    self.model.for_inference()
-    {indent}return output""",
-        )
+        new_code = indent + "if not _was_training:\n" + indent + "    self.model.for_inference()\n" + indent + "return output"
+        function = function.replace(f"{indent}return output", new_code)
 
     return function