ensure CUDAExecutionProvider is called in session when using nvidia gpu.

This commit is contained in:
catscarlet 2025-02-28 17:06:42 +08:00
parent 9079508935
commit aa6fb76e0f

View File

@ -13,9 +13,17 @@ class BaseSession:
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""Initialize an instance of the BaseSession class."""
self.model_name = model_name
device_type = ort.get_device()
if device_type == 'GPU' and 'CUDAExecutionProvider' in ort.get_available_providers():
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
else:
providers = ['CPUExecutionProvider']
self.inner_session = ort.InferenceSession(
str(self.__class__.download_models(*args, **kwargs)),
sess_options=sess_opts,
providers=providers,
)
def normalize(