mps support

This commit is contained in:
Wu Zichen
2024-01-24 19:37:47 +08:00
parent 8069264e64
commit 07a5339691
8 changed files with 70 additions and 33 deletions

View File

@@ -35,9 +35,11 @@ from my_utils import load_audio
from tools.i18n.i18n import I18nAuto
i18n = I18nAuto()
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # 确保直接启动推理UI时也能够设置。
if torch.cuda.is_available():
device = "cuda"
elif torch.mps.is_available():
elif torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"