【简单理解】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]])