How can I normalize after masking?

1

Suppose I have a tensor and a mask.

>>> a
tensor([[0.6000, 0.7000, 0.4000],
        [0.6000, 0.7000, 0.4000]])
>>> mask = torch.tensor([[1,0,0],[1,1,1]])
>>> a * mask
tensor([[0.6000, 0.0000, 0.0000],
        [0.6000, 0.7000, 0.4000]])
>>> F.softmax(a, 1)
tensor([[0.3420, 0.3780, 0.2800],
        [0.3420, 0.3780, 0.2800]])

However, I desire:


tensor([[1., 0., 0.],
        [0.3420, 0.3780, 0.2800]])
2

I got my answer in StackOverFlow: