Распараллелить простой цикл с использованием Pytorch на одном графическом процессоре

Question

У меня есть цикл for, который работает с независимыми столбцами большой матрицы. Я распараллелил цикл for на CPU, используя функцию prange в Numba. Теперь я хочу выполнить эту операцию, используя тензоры PyTorch на графическом процессоре. Я новичок в PyTorch и не знаю, как это сделать.

Любая помощь будет по достоинству оценена.

Мой код Python выглядит следующим образом:

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)

Мой код PyTorch выглядит следующим образом:

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)

Как я могу распараллелить циклы for?