CODEWITHSUNDEEP
pythonindexingpytorch
lima0
lima0

Reputation: 121

In a pytorch tensor, return an array of indices of the rows of specific value

Given the below tensor that has vectors of all zeros and vectors with ones and zeros:

tensor([[0., 0., 0., 0.],
    [0., 1., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [1., 0., 0., 1.],
    [0., 0., 0., 0.],...])

How can I have an array of indices of the vectors with ones and zeros so the output is like this:

indices = tensor([ 1, 3, 5, 6,...])

Update

A way to do it is:

indices = torch.unique(torch.nonzero(y>0,as_tuple=True)[0])

But I'm not sure if there's a better way to do it.

Upvotes: 1

Views: 803

Answers (1)

Ivan
Ivan

Reputation: 40618

An alternative way is to use torch.Tensor.any coupled with torch.Tensor.nonzero:

>>> x.any(1).nonzero()[:,0]
tensor([1, 3, 5, 6])

Otherwise, since the tensor contains only positive value, you can sum the columns and mask:

>>> x.sum(1).nonzero()[:,0]
tensor([1, 3, 5, 6])

Upvotes: 2

Related Questions