torch.nonzero

这个API传入一个Tensor参数，返回这个Tensor中所有不为0的下标。

>>> x=torch.tensor([[0,0,1],[1.,0,0],[1,1,1]])
tensor([[0., 0., 1.],
        [1., 0., 0.],
        [1., 1., 1.]])
        
>>> torch.nonzero(x)
tensor([[0, 2],
        [1, 0],
        [2, 0],
        [2, 1],
        [2, 2]])

这个API在GNN操作中还是蛮好用的，因为给定邻接矩阵，有时候我们需要对所有的边进行一个操作，使用这个API就可以定位出所有的边的下标。