I have a torch tensor of shape (batch_size, N). I want to apply functional softmax with dim 1 to this tensor, but I also want it to ignore zeros in the tensor and only apply it to non-zero values (the non-zeros in the tensor are positive numbers). I think what I am looking for is the sparse softmax.
I came up with this code: GitHub 1, but seems like it uses nn.Module instead of functional.
How can I use that sparse softmax as a functional one (that is, it is not a layer)?