Only auto enable bf16 VAE on nvidia GPUs that actually support it.

comfyanonymous · Jan 15, 2024 · f9e55d8 · f9e55d8
1 parent 2395ae7
commit f9e55d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -175,7 +175,7 @@ def is_nvidia():
         if int(torch_version[0]) >= 2:
             if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
                 ENABLE_PYTORCH_ATTENTION = True
-            if torch.cuda.is_bf16_supported():
+            if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8:
                 VAE_DTYPE = torch.bfloat16
     if is_intel_xpu():
         if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: