import torch


def gram_matrix(x):
    """ Calculates the Gram matrix for the
        feature maps contained in x.

    Parameters:
          x: feature maps
    Returns:
          G: gram matrix
    """
    b, c, h, w = x.size()
    F = x.view(b, c, h * w)
    G = torch.bmm(F, F.transpose(1, 2))
    G.div_(h * w)
    return G