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),)]