Do you need to use nn.MSELoss for your classification task?
If not, could you switch to nn.CrossEntropyLoss and remove the one-hot-encoding as the CrossEntropyLoss needs a target storing class indices.