cuda takes priority

hkchengrex · Dec 23, 2024 · d294fb8 · d294fb8
1 parent c8b5201
commit d294fb8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/demo.py b/demo.py
@@ -63,10 +63,10 @@ def main():
     mask_away_clip: bool = args.mask_away_clip
 
     device = 'cpu'
-    if torch.backends.mps.is_available():
-        device = 'mps'
-    elif torch.cuda.is_available():
+    if torch.cuda.is_available():
         device = 'cuda'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
     else:
         log.warning('CUDA/MPS are not available, running on CPU')
     dtype = torch.float32 if args.full_precision else torch.bfloat16

diff --git a/gradio_demo.py b/gradio_demo.py
@@ -19,10 +19,10 @@
 log = logging.getLogger()
 
 device = 'cpu'
-if torch.backends.mps.is_available():
-    device = 'mps'
-elif torch.cuda.is_available():
+if torch.cuda.is_available():
     device = 'cuda'
+elif torch.backends.mps.is_available():
+    device = 'mps'
 else:
     log.warning('CUDA/MPS are not available, running on CPU')
 dtype = torch.bfloat16