From 5a1b08bed54c6e99d4f80d2fc2562b65757b0f47 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 14 Jun 2022 15:19:38 +0200
Subject: [PATCH 1/2] Prefer MPS over CPU if available

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

diff --git a/hubconf.py b/hubconf.py
index 01f4eba08c81..df585f8cb411 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 d11df8337300..b329e3142c4f 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 hasattr(torch, 'has_mps') 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():

From b5b5d902c85eaa49edb995967f4722ffaf8dcf3b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 14 Jun 2022 15:30:46 +0200
Subject: [PATCH 2/2] Update torch_utils.py

---
 utils/torch_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/torch_utils.py b/utils/torch_utils.py
index b329e3142c4f..b1b107ee4f1b 100644
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -72,7 +72,7 @@ def select_device(device='', batch_size=0, newline=True):
             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
         arg = 'cuda:0'
-    elif not cpu and hasattr(torch, 'has_mps') and torch.backends.mps.is_available():  # prefer MPS if available
+    elif not cpu and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available():  # prefer MPS if available
         s += 'MPS\n'
         arg = 'mps'
     else:  # revert to CPU