diff --git a/extra/onnx_helpers.py b/extra/onnx_helpers.py index 7d3af8fa70..c11a8d41f8 100644 --- a/extra/onnx_helpers.py +++ b/extra/onnx_helpers.py @@ -1,10 +1,12 @@ from tinygrad import Tensor +from tinygrad.helpers import NUM_CPU_THREADS from tinygrad.tensor import _to_np_dtype from tinygrad.nn.onnx import OnnxRunner, OnnxValue import numpy as np import onnxruntime as ort ort_options = ort.SessionOptions() ort_options.log_severity_level = 3 +ort_options.intra_op_num_threads = NUM_CPU_THREADS.value def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}): """