Lua Torch相当于np.where()？

Question

我有一个,ByteTensor并希望获得有一个指数1.在numpy,我可以做类似的事情

a = np.array([1,0,1,0,1])
return np.where(a)

哪会回来(array([0, 2, 4]),).这个功能是在Torch中定义的吗？

(在我的特定情况下,我想使用这些索引来索引几个不同的Tensor对象,但是知道如何通常这样做会很好.)

Answer 1

你可以使用torch.nonzero,例如:

> a = torch.ByteTensor{1,0,1,0,1}
> print(torch.nonzero(a))                                                                                         
 1                                                                                                                  
 3                                                                                                                  
 5                                                                                                                  
[torch.LongTensor of size 3x1]

如果你真的需要找到1-s,你可以链接一个逻辑运算符:

> a = torch.ByteTensor{1,2,1,6,1}
> a:eq(1):nonzero()