From b754525e99ca62424c484fe529b6142f6bab939e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 12 Oct 2021 22:25:38 -0700 Subject: [PATCH] Check `'onnxruntime-gpu' if torch.has_cuda` (#5087) * Check `'onnxruntime-gpu' if torch.has_cuda` * fix indent --- detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/detect.py b/detect.py index 4e497305668c..d9961f5f457d 100644 --- 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