Lets say “a=torch.rand(1,4,2,2)” and “count_a=torch.zeros(4)”.
I want count of non zero of a(0,0,:,: ) to store in count_a(0), a(0,1,:,: ) to store in count_a(1) etc
If i used below code it stores total non zero count of “a” in to “count_a”.
Code:
import torch
count_a=torch.ones(4)
a=torch.rand(1,4,2,2)
count_a=torch.count_nonzero(a)
print("a value is ",a)
print("count is ",count_a)
Output:
a value is tensor([[[[0.8222, 0.4329],
[0.4672, 0.2243]],
     [[0.3086, 0.6668],
      [0.2699, 0.3848]],

     [[0.0251, 0.9675],
      [0.7624, 0.8663]],

     [[0.4905, 0.8086],
      [0.0784, 0.4926]]]])

count is tensor(16)
I want output to be tensor[4 4 4 4] as each layer has 4 non zero value instead of tesnor(16).