diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index e0e034fc5c..ed6207bb06 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -299,8 +299,9 @@ class FastMistralModel(FastLlamaModel):
             attention_module   = MistralAttention,
         )
         # Just for Mistral Nemo models!
-        function = patch_mistral_nemo_attention(function)
-        if True:#init_name is not None:
+        if function is not None:
+            function = patch_mistral_nemo_attention(function)
+            # if True:#init_name is not None:
             exec(function, globals())
             MistralAttention.__init__  = eval(init_name)
         pass