Prefer MPS over CPU if available (#8210)

* Prefer MPS over CPU if available

* Update torch_utils.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-06-17 18:28:29 +02:00
1 parent 86e6add07b
commit 303f3785ae
2 files changed
+8 -6

No files matched your search

+1 -1
View File
@@ -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