Error using MSELoss with CNN: Using a target size that is different to the input size


I am running a CNN experiment where my y shape is (32, 56, 56) and y_pred is (32, 313, 56, 56), 32 is batch size and images are 56x56. For classification when I use nn.CrossEntropyLoss() it accepts these tensors of different shapes but I also need to get results for MSELoss to show that MSE is not the better metric. I get following error:

loss_fn = nn.MSELoss(reduction='mean')

UserWarning: Using a target size (torch.Size([32, 56, 56])) that is different to the input size (torch.Size([32, 313, 56, 56])). 
This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
return F.mse_loss(input, target, reduction=self.reduction)

Can you please let me know how use MSELoss with different tensor shapes. I checked the parameters in nn.MSELoss() page but did not find anything I can use.

Thank You,

That’s true because the target is expected to contain class indices to index the class dimension of the model’s output tensor.

You won’t be able to directly compare class indices with logits which your model returns.
Instead you might need to one-hot encode the target tensor via F.one_hot and apply a softmax on your model output to create probabilities.
I would expect this use case to perform worse than the previous one using nnCrossEntropyLoss.