From 03ca52dc942930090859041483a855f72bf686f0 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 29 Jan 2024 03:43:39 +1100
Subject: [PATCH] saving

---
 unsloth/models/mistral.py | 3 +++
 unsloth/save.py           | 1 -
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 2941fb32fb..91a2dd7e35 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -375,6 +375,9 @@ class FastMistralModel(FastLlamaModel):
         
         # Log Unsloth version for future fastpaths for inference
         model.config.update({"unsloth_version" : __version__})
+
+        # Add save modules
+        patch_saving_functions(model)
         
         return model, tokenizer
     pass
diff --git a/unsloth/save.py b/unsloth/save.py
index baa8f3f5ca..4cda8b9651 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -278,7 +278,6 @@ def unsloth_save_model(
         not hasattr(internal_model.model, "layers")
     ):
         # Do general saving
-        print(type(model))
         # Edit save_pretrained_settings
         # [TODO] _create_repo has errors due to **kwargs getting accepted
         for deletion in \