pytorch

函数torch.gather(input, dim, index, out=None) → Tensor
沿给定轴 dim ,将输入索引张量 index 指定位置的值进行聚合.
对一个 3 维张量,输出可以定义为:

1
2
3

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

Parameters:

• input (Tensor) – 源张量
• dim (int) – 索引的轴
• index (LongTensor) – 聚合元素的下标(index需要是torch.longTensor类型)
• out (Tensor, optional) – 目标张量