diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index c55e4141d9..2d75452b20 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -71,7 +71,14 @@ def PatchRL(FastLanguageModel):
 pass
 
 
-# Handles NEFTune
+RLTrainer_replacement = '''
+import os
+from typing import *
+from dataclasses import dataclass, field
+from packaging.version import Version
+import torch
+
+# https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_utils.py#L126
 def neftune_post_forward_hook(module, input, output):
     if module.training:
         dims = torch.tensor(output.size(1) * output.size(2))
@@ -80,13 +87,6 @@ def neftune_post_forward_hook(module, input, output):
     return output
 pass
 
-
-RLTrainer_replacement = '''
-import os
-from typing import *
-from dataclasses import dataclass, field
-from packaging.version import Version
-
 @dataclass
 class Unsloth{RLConfig_name}({RLConfig_name}):
     """