[FIXED] How to create a mask for a tensor based on dimension-wise condition

Issue

Suppose I have a n by m tensor and I want to create a n by 1 binary mask so that mask[I]==0 if the row tensor[i] is all 0. How can I do this python pytorch?

Solution

If I'm reading the question correctly, you want something like this

x = torch.tensor([
    [0,1,2],
    [0,0,0],
    [0,1,0]
])

mask = x.eq(0).all(-1).long()
>tensor([0, 1, 0])

if you want to keep the unit axis, you can add keepdim=True

mask = x.eq(0).all(-1, keepdim=True).long()
>tensor([[0],
         [1],
         [0]])

Answered By - Karl