Jit compile custom C++ function with OpenMP threads


I am writing my custom layers in C++. Now I have created the Linear layer and I want to parallelize the GEMM function (I use a custom naive gemm). Can I jit compile with openmp threads? Because when I try to use openmp pragma for a parallel for, it seems pytorch give me weird errors that it does not know some of my Aten functions. i.e. (‘Tensor’ does not name a type; did you mean ‘THTensor’)

Thanks in advance