Based on the shape of y_train_pred it seems you are working on a multi-label classification?
If so, you shouldn’t use torch.argmax, as it would be used for a multi-class classification and would need to apply a threshold.
E.g. if the outputs are probabilities for each label and indicate, if this class index is active, you could use pred = output > 0.5.