@kimishpatel We’re using the python wheel for inference. I’ve tried recompiling pytorch, so that torch.set_num_threads calls pthreadpool()->set_thread_count as well (PR with the rebased commit 3). It works, but we have to ship a custom pytorch wheel. I didn’t find a way to do that through the python API or even with a C++ extension.