fix: manually set nb_threads for GPU

ELS-RD · Feb 17, 2023 · 8ddb0f4 · 8ddb0f4
1 parent f6dde83
commit 8ddb0f4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)