torch.gather()函数详解

input:

[[1,2,3],

[4,5,6],

[7,8,9]]

index:

[[2,1,0],

[1,0,1]]

如上已经给定了输入和索引，先假设dim=0

根据index，其本身元素的索引有（0,0）、（0,1）、（0,2）、（1,0）、（1,1）、（1,2）

因为设置dim=0，所以将index中的元素在dim=0这个维度进行替换，即得到（2,0）、（1,1）、（0,2）、（1,0）、（0,1）、（1,2），根据替换后的索引在input找值，得到输出

output：

[[7,5,3],

[4,2,6]]

如果设置dim=1，所以将index中的元素在dim=1这个维度进行替换，即得到（0,2）、（0,1）、（0,0）、（1,1）、（1,0）、（1,1），根据替换后的索引在input找值，得到输出

output:

[[3,2,1],

[5,4,5]]