u can use torch.nn.functional.softmax(input) to get the probability, then use topk function to get top k label and probability, there are 20 classes in your output, u can see 1x20 at the last line
btw, in topk there is a parameter named dimention to choose, u can get label or probabiltiy if u want
1 Like