From 16f413be6ced24bcc606ba9731c153f68d0a10b5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 7 Oct 2021 15:27:54 -0700 Subject: [PATCH 1/2] Check `'onnxruntime-gpu' if torch.has_cuda` --- detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/detect.py b/detect.py index 75ec3ecc5ff3..009326b0e8b1 100644 --- a/detect.py +++ b/detect.py @@ -86,7 +86,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s) modelc = load_classifier(name='resnet50', n=2) # initialize modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval() elif onnx: - 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 From be0d23640e44042e3dc7c2f24ec279ad111f3420 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 12 Oct 2021 22:03:18 -0700 Subject: [PATCH 2/2] fix indent --- detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/detect.py b/detect.py index ffdb8d8cf139..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-gpu' if torch.has_cuda else 'onnxruntime')) + check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime')) import onnxruntime session = onnxruntime.InferenceSession(w, None) else: # TensorFlow models