diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 39db053550..ed802e4877 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -19,6 +19,7 @@ __all__ = [
 ]
 
 import re
+import torch
 import inspect
 from collections import defaultdict
 RL_EXTRA_ARGS = defaultdict(list)