onnxruntime respects NUM_CPU_THREADS

2026-08-27 16:41:08 -07:00
1 changed files with 2 additions and 0 deletions
@@ -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={}):
  """