# Index into a pytorch tensor and return array of indices of rows

Sorry for the basic question, but I’m hoping the nice people here can help me:
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/rows with ones and zeros so the output is like this:

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

It looks like there is an implicit reduction here; are you treating rows with multiple ones the same as those that have a single one?
In that case, you could try:

``````import torch
a = torch.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.]])
print(a)
s = torch.sum(a, dim=-1)
print(s)
inds = torch.nonzero(s).reshape(-1)
print(inds)
``````
``````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.]])
tensor([0., 2., 0., 1., 0., 1., 2., 0.])
tensor([1, 3, 5, 6])
``````
1 Like