AI_Geeker的博客

原创 torch.gather函数的说明

torch.gather函数的理解 一.函数说明 torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor dim=0 时，output[i][j][k] = input[a][j][k] ， dim=1时，output[i][j][k] = input[i][a][k] dim=2时，output[i][j][k] = input[i][j][a] 上述input索引中的a=index[i][j][k]。 说明：out