【簡單理解】torch.nonzero（同樣適用於獲取tensor中某一元素的索引）

torch.nonzero其實就是找出tensor中非零的元素的索引

import torch
label = torch.tensor([[1,0,0],
                      [1,0,1]])
print(label.nonzero())

返回的結果：

tensor([[0, 0],
        [1, 0],
        [1, 2]])

返回的結果就是非零元素的索引，其中[0,0]對應了第一行第一列的1，[1,0]對應了第二行第一列的1，[1,2]對應了第二行第三列的1。

 

延伸：

有時我們只想得到一種元素對應的索引，比如我們只想要1對應的索引

import torch
label = torch.tensor([[1,0,0],
                      [3,0,1]])
print((label==1).nonzero())

返回的結果：

tensor([[0, 0],
        [1, 2]])

或者，我們想要一定條件下的元素的索引，比如大於1的元素的索引

import torch
label = torch.tensor([[1,0,0],
                      [3,0,1]])
print((label>1).nonzero())

返回結果：

tensor([[1, 0]])