glenn-jocher commited on
Commit
b754525
·
unverified ·
1 Parent(s): 20a809d

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

Browse files

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

* fix indent

Files changed (1) hide show
  1. detect.py +1 -1
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