add provider hook to run_pretrained_models (#1642)

Signed-off-by: Guenther Schmuelling <[email protected]>

Author: guschmue

tests/run_pretrained_models.py

@@ -335,14 +335,19 @@ def run_onnxruntime(self, name, model_proto, inputs, outputs, external_tensor_st
                                            as_text=utils.is_debug_mode(),
                                            external_tensor_storage=external_tensor_storage)
         logger.info("Model saved to %s", model_path)
+        providers = ['CPUExecutionProvider']
+        if rt.get_device() == "GPU":
+            gpus = os.environ.get("CUDA_VISIBLE_DEVICES")
+            if gpus is None or len(gpus) > 1:
+                providers = ['CUDAExecutionProvider']
+
         opt = rt.SessionOptions()
         if self.use_custom_ops:
             from onnxruntime_extensions import get_library_path
             opt.register_custom_ops_library(get_library_path())
-            m = rt.InferenceSession(model_path, opt)
         if self.ort_profile is not None:
             opt.enable_profiling = True
-        m = rt.InferenceSession(model_path, opt)
+        m = rt.InferenceSession(model_path, sess_options=opt, providers=providers)
         results = m.run(outputs, inputs)
         if self.perf:
             n = 0