Python 如何从二维pytorch张量得到最大元素的行和列索引？

是否有任何方法可以检索二维

pytorch

tensor中包含的最大元素的行和列索引？例如，请参见下面的

pytorch

tensor

a

：

a
>>火炬张量（[1,2,3]，
[9,5,4],
[6,7,8])

张量

a

中最大的元素是9，发生在第二行的第一列。如果我将其更改为从零开始的python列和行索引，元素的列索引将为0，行索引将为1

---

是否有任何方法可以从二维pytorch张量

a

中检索索引[1,0]？

不幸的是，没有内置方法。 但是，您可以使用numpy：

np.unravel_index(torch.argmax(a), a.shape)

否则，您需要编写自己的逻辑，例如：

def unravel_index(flat_idx, shape): 
     multi_idx = [] 
     r = flat_idx 
     for s in shape[:-1]: 
         multi_idx.append(r // s) 
         r = r % s 
     multi_idx.append(r % s) 
     return multi_idx