Normalization of multinomial probability

[LaurynasGiriunas]

First of all, thanks for a great paper and code! I wanted to ask if there is a way to normalize predicted values, the un-normalized multinomial probability?