import torch


def gram_matrix(input):
    """Computes the Gram matrix of a single image after a convolution layer"""
    a, b, c, d = input.size()  # a=batch size(=1)
                               # b=number of feature maps (convolution channels)
                               # (c,d)=dimensions of a f. map (N=c*d)
    features = input.view(a * b, c * d)  # resise F_XL into \hat F_XL

    G = torch.mm(features, features.t())  # compute the gram product

    # we 'normalize' the values of the gram matrix
    # by dividing by the number of element in each feature maps.
    # Note in Gatys paper this normalization is done in the loss, not here, but it's equivalent
    return G.div(a * b * c * d)