Sure! Here’s my code
import numpy as np
import torch
from torchvision import models

vgg = models.vgg19(pretrained=True)
modulelist = list(vgg.features.modules())
layer_inds = [0, 2]  # model layers to extract features from


def gram_matrix(activation, N, M):
    feature_map = activation.view(N, M)
    G = torch.mm(feature_map, feature_map.t())  # <-- Something's fishy here
    return G


def render(a, b, c):
    """Random render function, just want to generate some image..."""
    lin = torch.linspace(0, 1, 224)
    X, Y = torch.meshgrid(lin, lin)
    grid = torch.stack((X, Y, torch.zeros_like(X)))
    grid = grid * torch.stack((a, b, torch.tensor(1.0))).view(3, 1, 1)
    grid = (torch.sin(grid)) + 1.0 / 2.0

    return grid * c.view(3, 1, 1)


# Create ground truth
Pt = [torch.tensor(1.0), torch.tensor(50.), torch.tensor((0.5, 0.6, 0.7))]
truth = render(*Pt).unsqueeze(0)
w = torch.tensor([1e9, 1e9])


def gram_loss(img1, *params):
    img2 = render(*params).unsqueeze(0)
    E = []

    for i, layer in enumerate(modulelist[1:]):
        img1 = layer(img1)
        img2 = layer(img2)
        if i > layer_inds[-1]:
            break
        if i in layer_inds:
            N = img1.shape[1]
            M = img1.shape[2] * img1.shape[3]
            G1 = gram_matrix(img1, N, M)
            G2 = gram_matrix(img2, N, M)

            E.append(torch.sum(torch.pow(G1 - G2, 2)) * (1 / (4. * (N ** 2) * (M ** 2))))

    L_tot = (torch.stack(E) * w).sum()
    return L_tot


def run_gd():
    # Generate random starting parameters
    P = [torch.tensor(np.random.uniform(0, 1), requires_grad=True, dtype=torch.float32),
         torch.tensor(np.random.uniform(0, 50), requires_grad=True, dtype=torch.float32),
         torch.tensor(np.random.uniform(0, 1, size=(3,)), requires_grad=True, dtype=torch.float32)]

    optim = torch.optim.Adam(P, lr=0.01)

    for i in range(100):
        with torch.autograd.set_detect_anomaly(True):
            optim.zero_grad()
            loss = gram_loss(truth, *P)
            print("{}. Loss: {:.6f}".format(i + 1, loss.detach().numpy()))
            loss.backward()
            optim.step()


if __name__ == '__main__':
    run_gd()  # Run gradient descent

This is part of a bigger project to create a differentiable rendering process, hence the render() function and gradient descent. Also, this loss algorithm is based on this paper: https://arxiv.org/pdf/1505.07376.pdf 1