np.where用法

x = torch.tensor(
    [
        [1,1,0,0,0,1,1],
        [0,1,1,0,0,1,0],
        [0,1,0,0,0,1,0]
    ]
)
1. 整体一起 返回值（y[0][i],y[1][i]）构成坐标
y = np.where(x==1) 或者 y = np.where(x>0)

output:(array([0, 0, 0, 0, 1, 1, 1, 2, 2], dtype=int64),

            array([0, 1, 5, 6, 1, 2, 5, 1, 5], dtype=int64))

2. 一行一行来 返回每一行index

y = [np.where(i==1) for i in x]

output:[(array([0, 1, 5, 6], dtype=int64),),

            (array([1, 2, 5], dtype=int64),),

             (array([1, 5], dtype=int64),)]