saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-05 04:17:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[TensorRT EP] Avoid calling unavailable function with cpu python package (#19251)

Browse source 
C.register_tensorrt_plugins_as_custom_ops() is only available in gpu
python package.
Add condition to avoid calling it in cpu python package.
This commit is contained in:
Chi Lo 2024-01-23 21:30:22 -08:00 committed by GitHub
parent 6a424ccf8c
commit c10be1848c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
onnxruntime/python/onnxruntime_inference_collection.py
View file

@ -466,7 +466,7 @@ class InferenceSession(Session):
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)
@ -510,11 +510,15 @@ class InferenceSession(Session):
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])