理解torch.gather

先上案例：

>>> a = torch.Tensor([[1,2,3],[4,5,6]])
>>> a
tensor([[1., 2., 3.],
        [4., 5., 6.]])
>>> torch.gather(a,dim=0,torch.LongTensor([[0,1,0],[1,0,0]]))
tensor([[1., 5., 3.],
        [4., 2., 3.]])
>>> torch.gather(a,dim=1,torch.LongTensor([[2,0,1],[1,2,0]]))
tensor([[3., 1., 2.],
        [5., 6., 4.]])

a 为被索引的Tensor。
torch.LongTensor([[0,1,0],[1,0,0]])为索引。
dim = 0 : 固定索引的列号，然后确认行号。即，索引第一维度。
dim = 1 : 固定索引的行号，然后确认列号。即，索引第二维度。

如下图所示：