Check 'onnxruntime-gpu' if torch.has_cuda (#5087)
* Check `'onnxruntime-gpu' if torch.has_cuda` * fix indent
This commit is contained in:
1 parent
20a809dee5
commit
b754525e99
1 file changed
+1
-1
@@ -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
|
||||
|
||||
Reference in new issue
Block a user