Python 在Pytorch中改进索引以消除嵌套for循环

我有一个使用嵌套for循环的情况，但我想知道是否有一种更快的方法可以使用Pytorch中的一些高级索引来实现这一点

我有一个名为

t

的张量：

t = torch.randn(3,8)
print(t)
tensor([[-1.1258, -1.1524, -0.2506, -0.4339,  0.8487,  0.6920, -0.3160, -2.1152],
        [ 0.4681, -0.1577,  1.4437,  0.2660,  0.1665,  0.8744, -0.1435, -0.1116],
        [ 0.9318,  1.2590,  2.0050,  0.0537,  0.6181, -0.4128, -0.8411, -2.3160]])

我想创建一个新的张量，它对

t

中的值进行索引。 假设这些索引存储在变量

索引中

indexes = [[(0, 1, 4, 5), (0, 1, 6, 7), (4, 5, 6, 7)],
           [(2, 3, 4, 5)],
           [(4, 5, 6, 7), (2, 3, 6, 7)]]

索引中的每个内部元组表示从一行中获取的四个索引

例如，基于这些索引，我的输出将是一个6x4维张量（6是
索引中的元组总数，4对应于元组中的一个值）

例如，这就是我想要做的：

#counting the number of tuples in indexes
count_instances = sum([1 for lst in indexes for tupl in lst])

#creating a zero output matrix 
final_tensor = torch.zeros(count_instances,4)

final_tensor[0] = t[0,indexes[0][0]]
final_tensor[1] = t[0,indexes[0][1]]
final_tensor[2] = t[0,indexes[0][2]]
final_tensor[3] = t[1,indexes[1][0]]
final_tensor[4] = t[2,indexes[2][0]]
final_tensor[5] = t[2,indexes[2][1]]

最终输出如下所示：
打印（最终张量）

我创建了一个函数
build\u tensor
（如下所示）来通过嵌套for循环实现这一点，但我想知道在Pytorch中是否有一种通过简单索引来实现这一点的更快方法。我想要一种更快的方法，因为我用更大的索引和t大小做了数百次这个操作

有什么帮助吗

def build_tensor(indexes, t):
    #count tuples
    count_instances = sum([1 for lst in indexes for tupl in lst])
    #create a zero tensor
    final_tensor = torch.zeros(count_instances,4)
    final_tensor_idx = 0

    for curr_idx, lst in enumerate(indexes):
        for tupl in lst:
            final_tensor[final_tensor_idx] = t[curr_idx,tupl]
            final_tensor_idx+=1
    return final_tensor

您可以将索引排列成二维数组，然后一次完成索引，如下所示：

rows = [(row,)*len(index_tuple) for row, row_indices in enumerate(indexes) for index_tuple in row_indices]
columns = [index_tuple for row_indices in indexes for index_tuple in row_indices]
final_tensor = t[rows, columns]

这使得平均速度提高了5倍。谢谢
rows = [(row,)*len(index_tuple) for row, row_indices in enumerate(indexes) for index_tuple in row_indices]
columns = [index_tuple for row_indices in indexes for index_tuple in row_indices]
final_tensor = t[rows, columns]