import torch


def gram_matrix(input):
    """
    gets the gram matrix of the input and normalizes it
    :param input: a feature tensor of size (n, c, h, w)
    :return the gram matrix
    """
    n, c, h, w = input.size()
    features = input.view(n * c, h * w)
    gram = torch.mm(features, features.t())  # compute the Gram product

    # normalize Gram matrix
    return gram.div(n * c * h * w)