Lua等价于np.where()？

Question

我有一个ByteTensor，希望在有1的地方获取索引。在numpy中，我可以做一些类似的事情

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

这将返回(array([0, 2, 4]),)。这个功能是在火炬中定义的吗？

(在我的特殊情况下，我想使用这些索引来索引几个不同的张量对象，但是很高兴知道如何做到这一点。)

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，那么您可以链接一个逻辑运算符：

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