import torch


def gram_matrix(input_data):
    """ Gram Matrix for Style Loss Module

        a = batch size (1)
        b = number of feature maps
        c*d = number of features in a feature map

        Note: This specification is specific to 2d convolution
    
    """
    a, b, c, d = input_data.size()
    features = input_data.view(b, a * c * d) 
    G = torch.mm(features, features.t())
    return G.div(a * b * c * d)