Parallelize simple for loop using Pytorch on single GPU

I have a for-loop which operates on independent columns of a large matrix. I have parallelized the for-loop on CPU using the prange function in Numba. Now I want to perform this operation using PyTorch tensors on a GPU. I'm a PyTorch novice and don't know how to do it.

Any help will be really appreciated.

My Python code is as follows:

def select_next(X, gains, current_values, mask):
    for idx in prange(X.shape[0]):
        if mask[idx] == 1:
            continue

        a = numpy.maximum(X[idx], current_values)
        gains[idx] = (a - current_values).sum()
    return numpy.argmax(gains)

My PyTorch code is as follows:

def select_next(X, gains, current_values, mask):
    for idx in range(X.shape[0]):
        if mask[idx].item() == 1:
            continue

        a = torch.max(X[idx], current_values)
        gains[idx] = torch.sum(torch.sub(a, current_values))
    return torch.argmax(gains)

How can I parallelize the for loops?