gram_matrix.py

def gram_matrix(tensor):
    # Get the batch_size, depth, height, and width of the Tensor
    # Reshape it, so we're multiplying the features for each channel
    _, d, h, w = tensor.size()
    tensor = tensor.view(d, h*w)

    # Calculate the gram matrix by multiplying the tensor with its transpose
    gram = torch.mm(tensor, tensor.t())

    return gram