zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Check `'onnxruntime-gpu' if torch.has_cuda` (#5087) 

* Check `'onnxruntime-gpu' if torch.has_cuda`

* fix indent
This commit is contained in:
Glenn Jocher 2021-10-12 22:25:38 -07:00 committed by GitHub
parent 20a809dee5
commit b754525e99
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
detect.py
View File

@ -92,7 +92,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
# check_requirements(('opencv-python>=4.5.4',))
net = cv2.dnn.readNetFromONNX(w)
else:
check_requirements(('onnx', 'onnxruntime'))
check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime'))
import onnxruntime
session = onnxruntime.InferenceSession(w, None)
else: # TensorFlow models