Browse Source
Check `'onnxruntime-gpu' if torch.has_cuda` (#5087)
* Check `'onnxruntime-gpu' if torch.has_cuda` * fix indentmodifyDataloader
Glenn Jocher
GitHub
3 years ago
1 changed files with 1 additions and 1 deletions
+ 1
- 1
detect.py
View File
| # check_requirements(('opencv-python>=4.5.4',)) | # check_requirements(('opencv-python>=4.5.4',)) | ||||
| net = cv2.dnn.readNetFromONNX(w) | net = cv2.dnn.readNetFromONNX(w) | ||||
| else: | else: | ||||
| check_requirements(('onnx', 'onnxruntime')) | |||||
| check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime')) | |||||
| import onnxruntime | import onnxruntime | ||||
| session = onnxruntime.InferenceSession(w, None) | session = onnxruntime.InferenceSession(w, None) | ||||
| else: # TensorFlow models | else: # TensorFlow models |
Loading…