Merge pull request #4623 from fumiama/mps

Fix wrong mps fallback
This commit is contained in:
AUTOMATIC1111 2022-11-12 09:51:33 +03:00 committed by GitHub
commit 526f0aa556
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 3 deletions

View File

@ -3,8 +3,15 @@ import contextlib
import torch import torch
from modules import errors from modules import errors
# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility # has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
has_mps = getattr(torch, 'has_mps', False) # 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") cpu = torch.device("cpu")
@ -25,7 +32,7 @@ def get_optimal_device():
else: else:
return torch.device("cuda") return torch.device("cuda")
if has_mps: if has_mps():
return torch.device("mps") return torch.device("mps")
return cpu return cpu