Regards the tf.count_nonzero function from this page: https://www.tensorflow.org/api_docs/python/tf/count_nonzero 13
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 54 , seems like there is a way to count but it doesn’t keep the dimensions, how should I do this?
Thanks