From f0d2475e60fb3ab363d06ca0f81f65fec03faedb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 18 Mar 2025 01:44:50 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 1c7bdd0e6b..46069ce844 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -53,6 +53,7 @@ import math import functools from typing import Optional, Tuple, List, Union import re, inspect, sys +import contextlib import types try: from huggingface_hub.utils import get_token @@ -146,7 +147,11 @@ def unsloth_base_fast_generate( except: pass # Mixed precision autocast - with torch.inference_mode(), torch.autocast(device_type = "cuda", dtype = dtype): + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + autocaster = contextlib.nullcontext() + else: + autocaster = torch.autocast(device_type = "cuda", dtype = dtype) + with torch.inference_mode(), autocaster: try: output = self._old_generate(*args, **kwargs) except: