Exactly. You will need to perform the update step without no-grad mode for A to be L3’s graph, but this is a bit tricky since inplace updates (in grad-mode) aren’t allowed on parameters. You might want to use a library like higher, which handles this type of stuff for you. higher documentation — higher 0.2.1 documentation
import torch
import torch.nn as nn
import higher

model = nn.Linear(10, 10)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

x = torch.rand(10)
y = torch.rand(10, requires_grad=True)

with higher.innerloop_ctx(model, optimizer, copy_initial_weights=True) as (fmodel, diffopt):
    for i in range(2):
        loss = (fmodel(x) * y).sum()
        diffopt.step(loss)
        
        out = fmodel(x).sum()
        print(torch.autograd.grad(out, y))

Some other things to note:

you don’t need to do A.retain_grad() since A is already a leaf tensor
instead of doing A.grad = torch.autograd.grad(…, you can do loss.backward(inputs=A)

if you want the gradients to flow all the way back to A, you need to be careful of performing ops that autograd isn’t able to differentiate through, i.e., _, idx = torch.max(...).