From 44f3ab88aaec8c2bde6bf196adb9e8ecc21d0e3f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 18 Mar 2025 01:01:46 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index d79e9a829b..22057daf79 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -204,12 +204,11 @@ class FastBaseModel: if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 + elif os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + if dtype == torch.float16: dtype = torch.bfloat16 elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16: logger.warning_once("Device does not support bfloat16. Will change to float16.") dtype = torch.float16 - # Check forced float32 - if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - if dtype == torch.float16: dtype = torch.bfloat16 pass assert(dtype in (torch.float16, torch.bfloat16, torch.float32))