Solved by ptrblck in post #4 
              
              
                Thanks for the clarification! 
I misunderstood the [4, 4, 4, 4] output as the shapes, not the values. 
In that case, you could flatten the last two dimensions and use the dim argument in count_nonzero: 
a = torch.rand(1,4,2,2)
count_a = torch.count_nonzero(a.view(1, 4, -1), 2)
print("a value is ",a…