pytorch下角标索引

有两个形状一样的矩阵,然后怎么根据一个矩阵排好序得到的下标矩阵获取另一个矩阵呢?pytorch里面有这样的函数吗?
A=[[1,3,2]
      ,[4,3,1]]
B=[[1,2,3],
      [4,5,6]]
排序后下标相应为
[ [ 0,2,1],
   [2,1,0] ]
怎么由这些下标得到B'=
[[ B[0,0], B[0,2],B[0,1] ],
  [ B[1,2],B[1,1],B[1,0] ] ]
已邀请:

Mendel - 苦行僧

赞同来自:

已解决......
通过torch.gather可以实现

要回复问题请先登录注册