And here is the code for CUDA. I have only the forward for now:
std::vector<torch::Tensor> neighborfill_forward(torch::Tensor x, torch::Tensor y, torch::Tensor z, torch::Tensor neighbors){
    CHECK_INPUT(x);
    CHECK_INPUT(y);
    CHECK_INPUT(z);
    CHECK_CUDA(neighbors);
    
    return neighborfill_cuda_forward(x, y, z, neighbors)
}

std::vector<torch::Tensor> neighborfill_cuda_forward(torch::Tensor x, torch::Tensor y, torch::Tensor z, torch::Tensor neighbors){
    auto batch_size = x.size(0);
    auto state_size = x.size(2);
    int threads = 1024;
    dim3 blocks((state_size + threads - 1) / threads, batch_size);
    torch::Tensor nx = torch::zeros((1, x.size(1), neighbors.size(1) * neighbors.size(2)));
    torch::Tensor nx = torch::zeros((1, y.size(1), neighbors.size(1) * neighbors.size(2)));
    torch::Tensor nx = torch::zeros((1, z.size(1), neighbors.size(1) * neighbors.size(2)));
    
    neighborfill_cuda_forward_kernel<x.type(), neighbors.type()><<<blocks, threads>>>(
    nx.packed_accessor32<nx.type(), 3, torch::RestrictPtrTraits>(), 
    ny.packed_accessor32<ny.type(), 3, torch::RestrictPtrTraits>(), 
    nz.packed_accessor32<nz.type(), 3, torch::RestrictPtrTraits>(), 
    x.packed_accessor32<nx.type(), 3, torch::RestrictPtrTraits>(), 
    y.packed_accessor32<ny.type(), 3, torch::RestrictPtrTraits>(), 
    z.packed_accessor32<nz.type(), 3, torch::RestrictPtrTraits>(), 
    neighbors.packed_accessor32<neighbors.type(), 3, torch::RestrictPtrTraits>(), 
    x.size(2), neighbors.size(2), x.size(1)
    );
    
    return {nx, ny, nz}
}

template <typename x_t, typename neighbors_t>
__global__ void fill_neighbors_kernel(torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> nx, torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> ny, torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> nz, torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> x, torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> y, torch::PackedTensorAccessor32<x_t,3,torch::RestrictPtrTraits> z, torch::PackedTensorAccessor32<neighbor_t,3,torch::RestrictPtrTraits> neighbors, int tsize, int nsize, int fmaps){
    int index = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;
    if(index < tsize){ 
        for(int i = 0; i < nsize; i++){
            int nind = neighbors[0][index][i] - 1;   
            if(nind >= 0){
                for(int j = 0 ; j < fmaps; j++){
                    nx[0][j][index * nsize + i] = x[0][j][nind];//TODO INDEXING
                    ny[0][j][index * nsize + i]  = y[0][j][nind];
                    nz[0][j][index * nsize + i]  = z[0][j][nind];
                }
            }
        }
    }