在 Dim=1 的情况下应用 torch.gather()

上图显示了 torch gather() 函数在 dim=1 的二维张量上的工作。

这里索引张量的行对应于输入张量的行（用灰色阴影突出显示）。现在对于索引张量中的每个索引值，从该行和输入张量的索引中选取相应的值。

• 在第0行中，输入张量的第2、1、0个索引中的元素值分别为2、1和0。
• 在第 1 行中，输入张量的第 0、2 和 1 个索引中的元素值分别为 3、5 和 4。

torch.gather(input=tensor2, dim=1, index = torch.tensor([[2, 1, 0], 
                                                         [0, 2, 1]]))

张量（[[2, 1, 0],
        [3, 5, 4]])