Revert "Use onnxruntime-gpu"

This reverts commit 41d650f.
maxsitt · Nov 23, 2023 · 44ff3b9 · 44ff3b9
1 parent 5605068
commit 44ff3b9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/classify/predict.py b/classify/predict.py
@@ -397,7 +397,7 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop', 'onnxruntime'))
+    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
     run(**vars(opt))
 
 

diff --git a/classify/val.py b/classify/val.py
@@ -249,7 +249,7 @@ def parse_opt():
 
 
 def main(opt):
-    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop', 'onnxruntime'))
+    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
     run(**vars(opt))
 
 

diff --git a/export.py b/export.py
@@ -192,7 +192,7 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
     if simplify:
         try:
             cuda = torch.cuda.is_available()
-            check_requirements(('onnxruntime-gpu', 'onnx-simplifier>=0.4.1'))
+            check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
             import onnxsim
 
             LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')