Alternative for torch.distributions.categorical.Categorical()

Hello everbody,

I am novice in this field.

I am trying to replace torch.distributions.categorical.Categorical() with another method which will take an maximum value from softmax.

How can I deal with this situation?
Any help will be appreciated.