Hi
In my opinion, you do not need to create a mapping from index to label. Since your label is a tensor, just simply change your code to
for batch_idx, (data, label) in enumerate(train_loader_X):
    ...
    loss = loss_function(out, label - 1)
   ...

and
output = vae(test)
predicted = torch.argmax(output, dim=1) + 1