If “a=torch.rand(1,4,2,2)” then think of 4 layer having 2x2  value. If “a=torch.rand(1,8,2,2)” then think 8 layer having 2x2  value.
I want layer wise non zero count value. so for “a=torch.rand(1,4,2,2)” if all 4 layer 2x2 value is non zero then i want output to be “count_a=[4 4 4 4]”. Sorry i forgot to mention that count_a is of size “torch.Size([4])” and i want output to be tensor([4., 4., 4., 4.])
If “a=torch.rand(1,8,2,2)”,  if all the 8 layers 2x2 values are non zero then i want output to be tensor([4., 4., 4., 4., 4., 4., 4., 4. ]).