Can you give an example?
How would you write to this so it is correct?
    at::parallel_for(0, batch_size, 0, [&](int64_t start, int64_t end) {
        for (int64_t b = start; b < end; b++)
        {
            z_out[b] = z[b] - x[b];
        }
    });