在单个GPU上使用Pytorch并行化for循环

Question

我有一个for循环，可在大型矩阵的独立列上运行。我已经使用Numba中的prange函数并行化了CPU上的for循环。现在，我想在GPU上使用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循环？