please help me, as how can I use softmax in place of argmax() to get raw probability distribution for the predicted classes