b = torch.Tensor([[1,2,3],[4,5,6]])
print b
index_1 = torch.LongTensor([[0,1],[2,0]])
index_2 = torch.LongTensor([[0,1,1],[0,0,0]])
print torch.gather(b, dim=1, index=index_1)
print torch.gather(b, dim=0, index=index_2)
#结果
1 2 3
4 5 6
1 2
6 4
1 5 6
1 2 3
仔细研究一下结果是怎么来的, 对应原数组的索引位置
axis = 1
[ 1, 2 -> [ [0, 0], [0, 1]
6, 4 ] -> [1, 2], [1, 0] ]
结果就是在第0行中取第0, 1元素,在第1行取第2, 0元素
可以发现 axis = 1 结果 就是 自动添加行索引,在每一行中取值