pytorch中nonzero()的用法

pytorch中nonzero()的用法

nonzero（）：返回一个包含输入Input中非0元素索引的张量（以张量的形式返回非0索引）

import torch

a =torch.Tensor([1,2,4,4,5])
print(torch.nonzero(a<4))

a<4：输出的是布尔值索引（掩码）[1,1,0,0,0]
print(torch.nonzero(a<4))：输出布尔值索引非0的位置