import torch


def gram_matrix(x):
    """
    Arguments:
        x: a float tensor with shape [b, c, h, w].
    Returns:
        a float tensor with shape [b, c, c].
    """
    b, c, h, w = x.shape
    x = x.view(b, c, h * w)
    g = torch.matmul(x, x.permute(0, 2, 1))
    return g.div(c * h * w)