Solved by Kaii in post #5 
              
              
                x2.cpp 
#include <torch/extension.h>
#include <iostream>

#define _OPENMP
#include <ATen/ParallelOpenMP.h>

torch::Tensor x2(torch::Tensor z)
{
    torch::Tensor z_out = at::empty({z.size(0), z.size(1)}, z.options());
    int64_t batch_size = z.size(0);

    at::parallel_for(0, batch_size, 0, [&](in…