Extracting labels after applying softmax

Question

I have a multi class classification neural network. I apply softmax at the end to get probabilities for my classes. However, now I want to pick the maximum probability and get the corresponding label for it. I am able to extract the maximum probability but I'm confused how to get the label based on that. This is what I have:

labels = {'id1':0,'id2':2,'id3':1,'id4':3} ### labels

x_t = F.softmax(z,dim=-1)
#print(x_t)
y = torch.argmax(x_t, dim=1)
print(y.detach())

Can someone tell me how to get labels now. For example

y = tensor([3, 1, 3])

final_label = [id4,id3,id4]

Answer 1

You can try storing the map for index to label like this:

labels = {'id1':0,'id2':2,'id3':1,'id4':3} ### labels
idx_to_label = {v:k for k,v in labels.items()}

x_t = F.softmax(z,dim=-1)
#print(x_t)
y = torch.argmax(x_t, dim=1)
print(y.detach())
final_label = [idx_to_label[i] for i in y.detach().cpu().numpy()]
print(final_label)

Let me know if it helps.