def gram_matrices(features, normalize=True):
    """Calculate gram matrices (approximation of covariance matrix) for image features

    Args:
        features (tensor): shape NxCxHxW -- represents a batch of image features
        normalize (bool, optional): Whether or not to normalize the gram matrix. Defaults to True.

    Returns:
        tensor: gram matrices for each feature map
    """
    N, C, H, W = features.shape
    features = features.view(N, C, H * W)
    gram = features.bmm(features.permute(0, 2, 1)).contiguous()
    if normalize:
        gram /= (H * W * C)
    return gram