Merge pull request #4623 from fumiama/mps

Fix wrong mps fallback

1 files changed, 10 insertions, 3 deletions

diff --git a/modules/devices.py b/modules/devices.py
index 7511e1dc..bd3e4ffb 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -3,8 +3,15 @@ import contextlib
 import torch
 from modules import errors
 
-# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
-has_mps = getattr(torch, 'has_mps', False)
+# has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
+# check `getattr` and try it for compatibility
+def has_mps() -> bool:
+    if not getattr(torch, 'has_mps', False): return False
+    try:
+        torch.zeros(1).to(torch.device("mps"))
+        return True
+    except Exception:
+        return False
 
 cpu = torch.device("cpu")
 
@@ -25,7 +32,7 @@ def get_optimal_device():
         else:
             return torch.device("cuda")
 
-    if has_mps:
+    if has_mps():
         return torch.device("mps")
 
     return cpu