Hi, sorry that I did not read your code properly. _, y_pred_tags = torch.max(y_pred, dim = 1) does give the indices to you already.
From reading again, I am not too sure what’s the problem with your code, it seems good to me. Are u getting any errors from your run?