Add ONNX inference providers

Fix for #5916
ultralytics · Dec 8, 2021 · d80d80a · d80d80a
1 parent 554f782
commit d80d80a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/common.py b/models/common.py
@@ -322,7 +322,7 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False):
             LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
             check_requirements(('onnx', 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'))
             import onnxruntime
-            session = onnxruntime.InferenceSession(w, None)
+            session = onnxruntime.InferenceSession(w, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
         elif engine:  # TensorRT
             LOGGER.info(f'Loading {w} for TensorRT inference...')
             import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download