From 303f3785ae0fe8e8fff1eedebe25ec516a1baab1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 17 Jun 2022 18:28:29 +0200
Subject: [PATCH] Prefer MPS over CPU if available (#8210)

* Prefer MPS over CPU if available

* Update torch_utils.py
---
 hubconf.py           |  2 +-
 utils/torch_utils.py | 12 +++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/hubconf.py b/hubconf.py
index 01f4eba..df585f8 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -41,7 +41,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
     name = Path(name)
     path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name  # checkpoint path
     try:
-        device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
+        device = select_device(device)
 
         if pretrained and channels == 3 and classes == 80:
             model = DetectMultiBackend(path, device=device)  # download/load FP32 model
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
index d11df83..b1b107e 100644
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -62,8 +62,7 @@ def select_device(device='', batch_size=0, newline=True):
         assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
             f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
 
-    cuda = not cpu and torch.cuda.is_available()
-    if cuda:
+    if not cpu and torch.cuda.is_available():  # prefer GPU if available
         devices = device.split(',') if device else '0'  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
         n = len(devices)  # device count
         if n > 1 and batch_size > 0:  # check batch_size is divisible by device_count
@@ -72,15 +71,18 @@ def select_device(device='', batch_size=0, newline=True):
         for i, d in enumerate(devices):
             p = torch.cuda.get_device_properties(i)
             s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB
-    elif mps:
+        arg = 'cuda:0'
+    elif not cpu and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available():  # prefer MPS if available
         s += 'MPS\n'
-    else:
+        arg = 'mps'
+    else:  # revert to CPU
         s += 'CPU\n'
+        arg = 'cpu'
 
     if not newline:
         s = s.rstrip()
     LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)  # emoji-safe
-    return torch.device('cuda:0' if cuda else 'mps' if mps else 'cpu')
+    return torch.device(arg)
 
 
 def time_sync():