Count nonzero elements across dimension an keep dimensions

1

Regards the tf.count_nonzero function from this page: https://www.tensorflow.org/api_docs/python/tf/count_nonzero

tf.count_nonzero(x, 1, keepdims=True)  # [[1], [2]]

Is there a way that I can do similar things in Pytorch and still keep the dimension?
From the answer in this thread Count nonzeros element along an axis , seems like there is a way to count but it doesn’t keep the dimensions, how should I do this?

Thanks

2

Hi,

You can give the keepdim=True argument to the sum?

1 Like
3

I see, thanks. It’s the solution :slight_smile: