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, 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)?