import torch


def nchw_l2(x, y, squared=True):
    """ Computes l2 norm between two NxCxHxW images
    ARGS:
        x, y: Tensor/Variable (NxCxHxW) - x, y must be same type & shape.
        squared : bool - if True we return squared loss, otherwise we return
                         square root of l2
    RETURNS:
        ||x - y ||_2 ^2 (no exponent if squared == False),
        shape is (Nx1x1x1)
    """
    temp = torch.pow(x - y, 2) # square diff


    for i in range(1, temp.dim()): # reduce on all but first dimension
        temp = torch.sum(temp, i, keepdim=True)

    if not squared:
        temp = torch.pow(temp, 0.5)

    return temp.squeeze()