Skip to content

Commit

Permalink
[TensorRT EP] Avoid calling unavailable function with cpu python pack…
Browse files Browse the repository at this point in the history
…age (#19251)

C.register_tensorrt_plugins_as_custom_ops() is only available in gpu
python package.
Add condition to avoid calling it in cpu python package.
  • Loading branch information
chilo-ms authored Jan 24, 2024
1 parent 6a424cc commit c10be18
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions onnxruntime/python/onnxruntime_inference_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -466,7 +466,7 @@ def _create_inference_session(self, providers, provider_options, disabled_optimi

session_options = self._sess_options if self._sess_options else C.get_default_session_options()

self._register_ep_custom_ops(session_options, providers, provider_options)
self._register_ep_custom_ops(session_options, providers, provider_options, available_providers)

if self._model_path:
sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
Expand Down Expand Up @@ -510,11 +510,15 @@ def _reset_session(self, providers, provider_options):
self._sess_options = self._sess_options_initial
self._create_inference_session(providers, provider_options)

def _register_ep_custom_ops(self, session_options, providers, provider_options):
def _register_ep_custom_ops(self, session_options, providers, provider_options, available_providers):
for i in range(len(providers)):
if providers[i] == "TensorrtExecutionProvider":
if providers[i] in available_providers and providers[i] == "TensorrtExecutionProvider":
C.register_tensorrt_plugins_as_custom_ops(session_options, provider_options[i])
elif isinstance(providers[i], tuple) and providers[i][0] == "TensorrtExecutionProvider":
elif (
isinstance(providers[i], tuple)
and providers[i][0] in available_providers
and providers[i][0] == "TensorrtExecutionProvider"
):
C.register_tensorrt_plugins_as_custom_ops(session_options, providers[i][1])


Expand Down

0 comments on commit c10be18

Please sign in to comment.