How can I get the MSE of a tensor across a specific dimension?

Question

I have 2 tensors with .size of torch.Size([2272, 161]). I want to get mean-squared-error between them. However, I want it along each of the 161 channels, so that my error tensor has a .size of torch.Size([161]). How can I accomplish this?

It seems that torch.nn.MSELoss doesn't let me specify a dimension.

Answer 1

For the nn.MSELoss you can specify the option reduction='none'. This then gives you back the squared error for each entry position of both of your tensors. Then you can apply torch.sum/torch.mean.

a = torch.randn(2272,161)
b = torch.randn(2272,161)
loss = nn.MSELoss(reduction='none')
loss_result = torch.sum(loss(a,b),dim=0) 

I don't think there is a direct way to specify at the initialisation of the loss to which dimension to apply mean/sum. Hope that helps!