Solved by ptrblck in post #2 
              
              
                I think you could use a single output unit in the last layer via: 
model.fc = nn.Linear(fc_in_features, 1)

and use nn.MSELoss directly by providing the target in the shape [batch_size, 1] containing the class indices. 
As you said, you would usually use e.g. nn.CrossEntropyLoss, but if your target …