Hi, I’m trying to differentiate through an optimization procedure of another neural net. The code below shows 2 functions:

a working version where instead of neural nets we just have single tensors
the version with neural nets with 3 attempted solutions commented out.

Both should be equivalent.
I’m a newbie so I welcome explanations of what I’m actually doing in those failed attempts.
For context, I’m implementing the LOLA paper where agents in games model each other’s learning: https://arxiv.org/pdf/1709.04326.pdf 7
import torch

gd_iter = 5
lr = 0.1
A = torch.tensor([[-2., -3.], [0., -1.]])


class SelfOutputNet(torch.nn.Module):
   """Net that just returns it's own parameter vector"""
    def __init__(self):
        super(SelfOutputNet, self).__init__()
        self.theta = torch.nn.Parameter(torch.zeros(2, requires_grad=True))
    def forward(self):
        return self.theta


def test1_tensor_working():
   """Uses tensors"""
    theta1 = torch.zeros(2, requires_grad=True)
    theta2 = torch.zeros(2, requires_grad=True)
    theta1_optimizer = torch.optim.SGD((theta1,), lr=lr)

    def fake_objective_func(x1, x2):
        """Prisoner's dilemma loss"""
        x1, x2 = torch.softmax(x1, 0), torch.softmax(x2, 0)
        return x1.view(1, -1).mm(A).view(-1).dot(x2)

    # Inner optimization
    for k in range(gd_iter):
        fake_objective2 = fake_objective_func(theta2, theta1)
        grad2 = torch.autograd.grad(fake_objective2, theta2, create_graph=True)[0]
        theta2 = theta2 - lr * grad2  # Produces the result it should with my true objective func
   
   # Outer optimization
    fake_objective1 = fake_objective_func(theta1, theta2)
    theta1_optimizer.zero_grad()
    fake_objective1.backward()
    theta1_optimizer.step()


def test2_NN_not_working():
   """Equivalent test using nn.Modules"""
    net1 = SelfOutputNet()
    net2 = SelfOutputNet()
    theta1_optimizer = torch.optim.SGD(net1.parameters(), lr=lr)

    def fake_objective_func(x1, x2):
        """Prisoner's dilemma loss"""
        out1, out2 = x1.forward(), x2.forward()
        x1, x2 = torch.softmax(out1, 0), torch.softmax(out2, 0)
        return - x1.view(1, -1).mm(A).view(-1).dot(x2)

    # Inner optimization
    for k in range(gd_iter):
        fake_objective2 = fake_objective_func(net2, net1)
        grad2 = torch.autograd.grad(fake_objective2, net2.parameters(), create_graph=True)
        for i, layer in enumerate(net2.parameters()):
            """Three failing solutions:"""
            # layer.sub_(lr * grad2[i])       # RuntimeError: a leaf Variable that requires grad has been used in an in-place operation.
            # layer.data.sub_(lr * grad2[i])  # Has no effect as measured by final outcome when using my true objective func
            # layer = lr * grad2[i]           # Has no effect

   # Outer optimization
    fake_objective1 = fake_objective_func(net1, net2)
    theta1_optimizer.zero_grad()
    fake_objective1.backward()
    theta1_optimizer.step()

test_tensor_working()
test_NN_not_working()```