Check `'onnxruntime-gpu' if torch.has_cuda` (#5087)
Browse files* Check `'onnxruntime-gpu' if torch.has_cuda`
* fix indent
detect.py
CHANGED
@@ -92,7 +92,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
92 |
# check_requirements(('opencv-python>=4.5.4',))
|
93 |
net = cv2.dnn.readNetFromONNX(w)
|
94 |
else:
|
95 |
-
check_requirements(('onnx', 'onnxruntime'))
|
96 |
import onnxruntime
|
97 |
session = onnxruntime.InferenceSession(w, None)
|
98 |
else: # TensorFlow models
|
|
|
92 |
# check_requirements(('opencv-python>=4.5.4',))
|
93 |
net = cv2.dnn.readNetFromONNX(w)
|
94 |
else:
|
95 |
+
check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime'))
|
96 |
import onnxruntime
|
97 |
session = onnxruntime.InferenceSession(w, None)
|
98 |
else: # TensorFlow models
|