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

* Check `'onnxruntime-gpu' if torch.has_cuda` * fix indent
fhkiel-mlaip · Oct 13, 2021 · 2ad33cc · 2ad33cc
1 parent 1384355
commit 2ad33cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/detect.py b/detect.py
@@ -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