diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 46069ce844..acf999faff 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -148,7 +148,7 @@ def unsloth_base_fast_generate( # Mixed precision autocast if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - autocaster = contextlib.nullcontext() + autocaster = torch.autocast(device_type = "cuda", dtype = dtype) else: autocaster = torch.autocast(device_type = "cuda", dtype = dtype) with torch.inference_mode(), autocaster: