Add warning when inference with mps and bf16 on Mac

Omyzmw · Jun 4, 2024 · 1c5fadf · 1c5fadf
1 parent dfbc321
commit 1c5fadf
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/web_demo.py b/web_demo.py
@@ -26,7 +26,11 @@
 device = args.device
 assert device in ['cuda', 'mps']
 if args.dtype == 'bf16':
-    dtype = torch.bfloat16
+    if device == 'mps':
+        print('Warning: MPS does not support bf16, will use fp16 instead')
+        dtype = torch.float16
+    else:
+        dtype = torch.bfloat16
 else:
     dtype = torch.float16