Delete all zero'ed rows from a tensor

Hello everyone!

I have a tensor, let’s say A, which contains zero rows. I mean rows with all zero elements. There are other rows with non-zero elements as well. I there any elegant way to remove all zero’ed rows? To shed a light on the question, suppose the following example.

A=
[ [ [0,0,0,0,0], [0,-0.4,0,-0.9,0], [0,-0.7,0.9,0.2,0], [0,0,0,0,0] ] ]

What I want to get (let’s say tensor B) is the following:

B=
[ [ [0,-0.4,0,-0.9,0], [0,-0.7,0.9,0.2,0] ] ]

A = torch.Tensor([ [ [0,0,0,0,0], [0,-0.4,0,-0.9,0], [0,-0.7,0.9,0.2,0], [0,0,0,0,0] ] ])
B = A[A.sum(dim=2) != 0]
tensor([[ 0.0000, -0.4000, 0.0000, -0.9000, 0.0000],
[ 0.0000, -0.7000, 0.9000, 0.2000, 0.0000]])

Hi Sunny (and Sajad)!

That’s a nice solution.

Note, you should add an absolute-value in there to avoid deleting
a non-zero row that sums to zero, such as [1.0, -1.0, 0.0].

Best.

K. Frank

1 Like

Yeah, that’s right. I forgot it.