googlebot:

I’ll paste just in case


I updated the function inspired by your code. It takes 1.5x longer than the unoptimized version to run on my PC. I don’t think it is being vectorized.
torch::Tensor f_forward_vec(torch::Tensor input){
    const torch::Tensor& output = torch::zeros_like(input);
    const int batch_stride = input.strides()[0]; 
    const int half_batch_stride = batch_stride / 2;
    const int onehalf_batch_stride = batch_stride + half_batch_stride;
    float *__restrict__ in_i1 = input.data_ptr<float>();
    float *__restrict__ in_i2 = in_i1 + half_batch_stride;
    float *__restrict__ out_i1 = output.data_ptr<float>();
    float *__restrict__ out_i2 = out_i1 + half_batch_stride;
    const float* end = input.data_ptr<float>() + input.numel();
      
    // This loop goes over samples in the batch
    while(in_i2 < end){ 
        // This loop goes over elements in a sample
        for (int i=0; i<batch_stride; i++){ 
            float v1 = in_i1[i];
            float v2 = in_i2[i];
            bool b1 = std::signbit(v1), b2 = std::signbit(v2);
            out_i1[i] = b1 ? (b2 ? v2 : v1+v2) : (b2 ? v1+v2 : v1);
            out_i2[i] = b1 ? (b2 ? v1 : v1+v2) : (b2 ? v1+v2 : v2);      
        }
        in_i1 += onehalf_batch_stride;
        in_i2 += onehalf_batch_stride;
        out_i1 += onehalf_batch_stride;
        out_i2 += onehalf_batch_stride;
    }    
    return output;
}

It’s no big deal though. I care about the Cuda version which turned out to work great.