add mps support

Cookerjin · May 30, 2024 · b978e88 · b978e88
1 parent 934cbe3
commit b978e88
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/ChatTTS/utils/gpu_utils.py b/ChatTTS/utils/gpu_utils.py
@@ -16,6 +16,8 @@ def select_device(min_memory = 2048):
         if free_memory_mb < min_memory:
             logger.log(logging.WARNING, f'GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left.')
             device = torch.device('cpu')
+    elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
+        device = torch.device('mps')
     else:
         logger.log(logging.WARNING, f'No GPU found, use CPU instead')
         device = torch.device('cpu')