torch.scatter

https://zhuanlan.zhihu.com/p/426906301

 

 self[index[i][j][k]][j][k] = src[i][j][k]  # if dim == 0
 self[i][index[i][j][k]][k] = src[i][j][k]  # if dim == 1
 self[i][j][index[i][j][k]] = src[i][j][k]  # if dim == 2