import torch


def mse_loss(x, x_hat):
    """
    Returns the MSE between an image and its reconstruction

    INPUT:
    x: Tensor (B, C, H, W) -> source image
    x_hat: Tensor (B, C, H, W) -> reconstruction

    OUTPUT:
    mse_loss_sum: Tensor (B, 1) -> sum of the MSE for each image pair
    """
    mse_loss = torch.nn.MSELoss(reduction='none')
    recons_error = mse_loss(x, x_hat)
    mse_loss_sum = torch.sum(recons_error, dim=(1, 2, 3))

    return mse_loss_sum