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