我有一个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循环?