From b4d0680de48c3e4347a36733baa2bf14cb7e73f9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 15 May 2025 02:17:03 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a712f4a40b..94e2b4c3d4 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -313,12 +313,14 @@ class FastBaseModel: # Check for custom data-types custom_datatype = None + correct_dtype = 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) + correct_dtype = bnb_compute_dtype pass # Stop SDPA for some archs like Pixtral / Mistral3 @@ -432,6 +434,7 @@ class FastBaseModel: downcast_rope = False, fix_embeddings = False, do_forced_float32 = do_forced_float32, + correct_dtype = correct_dtype, ) model, tokenizer = patch_tokenizer(model, tokenizer) model = post_patch_loss_function(model)