From 3a9c10c33ce40020743735c0d92e85f9e554c6bb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 15 May 2025 01:45:15 -0700
Subject: [PATCH] Fix Seasame

---
 unsloth/models/loader.py |  2 +-
 unsloth/models/vision.py | 18 +++++++++++-------
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 1099c4f0de..8547102a25 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -543,7 +543,7 @@ class FastModel(FastBaseModel):
         elif "csm-1b" in model_name.lower():
             os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
             os.environ["UNSLOTH_DISABLE_FAST_GENERATION"] = "1"
-            os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
+            os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.floa16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
         elif "olmo-2" in model_name.lower() and transformers_version < Version("4.50.0.dev0"):
             raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)
         pass
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 48e8b532f2..9005150a50 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -303,13 +303,6 @@ class FastBaseModel:
         pass
         assert(dtype in (torch.float16, torch.bfloat16, torch.float32))
 
-        # Check for custom data-types
-        custom_datatype = None
-        if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "":
-            custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"]
-            dtype = torch.float32
-        pass
-
         bnb_compute_dtype = dtype
         do_forced_float32 = False
         if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
@@ -317,6 +310,17 @@ class FastBaseModel:
             bnb_compute_dtype = torch.float16
             do_forced_float32 = True
         pass
+
+        # Check for custom data-types
+        custom_datatype = None
+        if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "":
+            custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"]
+            assert custom_datatype.count(";") == 1
+            bnb_compute_dtype, custom_datatype = custom_datatype.split(";", 1)
+            dtype = torch.float32
+            bnb_compute_dtype = eval(bnb_compute_dtype)
+        pass
+
         # Stop SDPA for some archs like Pixtral / Mistral3
         if not ("attn_implementation" in kwargs):
             kwargs["attn_implementation"] = "sdpa"