Pytorch has torch.nn.MSELoss() and the functional torch.nn.functional.mse_loss().
Passing reduce=False to any of these will keep your dimensions:
a = torch.rand([2, 3, 20, 20])
b = torch.rand([2, 3, 20, 20])

torch.nn.functional.mse_loss(a, b, reduce=False).shape

> torch.Size([2, 3, 20, 20])