From 8ddb0f4ee30a0bc7b2f248a9a93beafe19c12bb9 Mon Sep 17 00:00:00 2001 From: Ferdinand Mom Date: Fri, 17 Feb 2023 14:33:30 +0000 Subject: [PATCH] fix: manually set nb_threads for GPU --- src/transformer_deploy/backends/ort_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformer_deploy/backends/ort_utils.py b/src/transformer_deploy/backends/ort_utils.py index 939f8f5a..9a9f3a63 100644 --- a/src/transformer_deploy/backends/ort_utils.py +++ b/src/transformer_deploy/backends/ort_utils.py @@ -81,7 +81,7 @@ def create_model_for_provider( provider_to_use = [provider_to_use] if provider_to_use == ["CPUExecutionProvider"]: options.execution_mode = ExecutionMode.ORT_SEQUENTIAL - options.intra_op_num_threads = nb_threads + options.intra_op_num_threads = nb_threads return InferenceSession(path, options, providers=provider_to_use)