I have tensor of size [3,1,2,2] ( Consider 3 batch of size (1,2,2) } . I want number which is greater than 0.1 in each batch. Below code gives total count greater than 0.1 which is 8. I want array of [4,3,1] corresponding to value greater than 0.1 in each batch.

What should the output tensor contain for the values smaller than your threshold? If you would like to fill it with default value, you could initialize the result tensor with these default values and use scatter to fill in the values meeting your threshold (torch.where might also work).

In below tensor a[0] has 4 value which is >0.1, a[1] has 3 value which is >0.1 and a[2] has 1 value which is >0.1. The output is giving sum of (4+3+1=8). I want tensor of [4,3,1].

A tensor in the shape [4, 3, 1] has 3 dimensions and contains 12 elements (4*3*1).
Print this random tensor in the same shape to see how the values are stored: print(torch.randn(4, 3, 1)).
Since your output has 8 values, you won’t be able to use a tensor with 12 elements to store them (unless you want to use default values for the other indices as described in the previous post).

Ah, it seems I’ve misunderstood the question and you want to get the number of elements > 0.1 in dim=[2, 3]?
If so, this should work: b = torch.sum(a>0.1, dim=[2, 3]).