PyTorch Hub models default to CUDA:0 if available (#2472)

* PyTorch Hub models default to CUDA:0 if available

* device as string bug fix
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-03-14 23:16:17 -07:00
1 parent 2d41e70e82
commit 9b11f0c58b
4 files changed
+9 -7

No files matched your search

+3 -1
View File
@@ -12,6 +12,7 @@ import torch
from models.yolo import Model
from utils.general import set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device
dependencies = ['torch', 'yaml']
set_logging()
@@ -43,7 +44,8 @@ def create(name, pretrained, channels, classes, autoshape):
model.names = ckpt['model'].names # set class names attribute
if autoshape:
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
return model
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
return model.to(device)
except Exception as e:
help_url = 'https://github.com/ultralytics/yolov5/issues/36'