torch.gather函数的说明

说明：index的索引与output的索引相同。即output[0][1][3],index的索引也是[0][1][3]
二.例子
input=torch.randn((2,2,2))
tensor([[[-1.0419, -0.6704],
[ 0.4355, 0.3973]],
[[-0.6508, -0.3114],
[-0.3993, -1.0891]]])
index=torch.LongTensor([[[1, 0],
[1, 1]],
[[1, 0],
[0, 1]]])

output=torch.gather(input,0,index)
dim=0,按照我们前面的解释，
输出output[0][0][0]的元素应该是input[a][0][0]的元素，其中a=index[0][0][0]=1
故ouput[0][0][0]=input[1][0][0]=-0.6508
输出output[1][0][1]处的元素应该等于input[a][0][1],其中a=index[1][0][1]=0
output[1][0][1]=input[0][0][1]=-0.6704
其他位置的元素依次类推。
output
tensor([[[-0.6508, -0.6704],
[-0.3993, -1.0891]],
[[-0.6508, -0.6704],
[ 0.4355, -1.0891]]])