From 69b09d80c800a0a072cf39992f602920efad47d7 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Jan 2025 01:22:13 -0800
Subject: [PATCH] Fix

---
 unsloth/models/_utils.py  | 10 ++++------
 unsloth/models/mistral.py |  2 --
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 279064b5eb..ff2c8726e4 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -799,9 +799,7 @@ def patch_linear_scaling(
         f"from typing import Union, Optional, List, Any, Callable, Tuple\n"\
         f"from {model_filepath} import logger, "\
         f"{model_name.title()}Attention, {model_name.title()}Config"
-
-    print(exec_code)
-    print(inspect.getsource(attention_module.__init__))
+    
     try:
         function = inspect.getsource(attention_module.__init__)
     except:
@@ -845,12 +843,12 @@ def patch_linear_scaling(
         "self.rotary_emb = .+?\)", function,
         flags = re.DOTALL | re.MULTILINE,
     )
-    if len(rotary_emb) == 0: return None, function
+    if len(rotary_emb) == 0:
+        return None, exec_code + "\n\n" + function
+
     rotary_emb = rotary_emb[0]
     function = function.replace(rotary_emb, fix_rope_function, 1)
     function = exec_code + "\n\n" + function
-    print("###########")
-    print(exec_code)
     return init_name, function
 pass
 
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 4edc3b7990..9a97015f9b 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -305,9 +305,7 @@ class FastMistralModel(FastLlamaModel):
         )
         # Just for Mistral Nemo models!
         if function is not None:
-            print(function)
             function = patch_mistral_nemo_attention(function)
-            print(function)
             # if True:#init_name is not None:
             exec(function, globals())
             MistralAttention.__init__  = eval(init_name)