Skip to content

Commit 76ab31e

Browse files
committed
Fix wrong mps selection below MasOS 12.3
1 parent 7ba3923 commit 76ab31e

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

modules/devices.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,15 @@
33
import torch
44
from modules import errors
55

6-
# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
7-
has_mps = getattr(torch, 'has_mps', False)
6+
# has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
7+
# check `getattr` and try it for compatibility
8+
def has_mps() -> bool:
9+
if getattr(torch, 'has_mps', False): return False
10+
try:
11+
torch.zeros(1).to(torch.device("mps"))
12+
return True
13+
except Exception:
14+
return False
815

916
cpu = torch.device("cpu")
1017

@@ -25,7 +32,7 @@ def get_optimal_device():
2532
else:
2633
return torch.device("cuda")
2734

28-
if has_mps:
35+
if has_mps():
2936
return torch.device("mps")
3037

3138
return cpu

0 commit comments

Comments
 (0)