How can I remove zero'd elements from a tensor?

I have a variable that looks something like [1,2,3,0,0,0,1,2] and I want to end up with the following variable: [1,2,3,1,2].

In TensorFlow, this would be done with something along the lines of tf.gather(tf.where(tf.not_equal(a, zero_tensor))). Is there a way to do this in PyTorch?

Thanks in advance.

Hi,

If your tensor is a, you can simply do: a[a.nonzero()].

7 Likes

I tried that but got:
autograd/_functions/tensor.py", line 328, in forward
assert not ctx.needs_input_grad[2]
AssertionError

Are you sure that this works on variables?

It works with pytorch 0.4,
If your using 0.3, you migh want to do: a[a.nonzero().detach()]: you cannot backpropagate the indices so they should not require gradients. I don’t have pytorch 0.3 here, does that work?

Now it’s giving the error:
return tensor.index_select(dim, index)
RuntimeError: invalid argument 3: expecting vector of indices at

a[a.nonzero().squeeze().detach()] ? :slight_smile:

1 Like

Still not working :frowning:
Both a.nonzero() and a.nonzero() are torch.cuda.LongTensor of size 10x2, but for some reason only the latter throws the above error.

By the way, is PyTorch 0.4 only for Python 3.5? I upgraded my PyTorch but it only installed 3.1…

OK, I think I solved it using detach_() instead of detach()

Pytorch 0.4 is not released yet. It is still under development and partially available if you compile from master.

Ok I think the problem is that your input a is 2D of size (1, something) and thus nonzero() returns a tensor with 2 coordinates with one which is always 0. If you make your input a 1D tensor, then nonzero() will return a 10x1 tensor which will be transformed into a vector of size 10 which you can then use to index your tensor.

For the detach(), I’m not sure why you would want to use the inplace version. The regular version should work here.