import torch


def get_sse(x_hat, x):
    """
    Calculates ||x_hat - x||^2 for each image in the batch dimension
    """
    sse_val = torch.sum((x_hat - x)**2, dim=[1,2,3])

    return sse_val.cpu().numpy().flatten()