AlbanD, this does quite work. I wanted to make sure it worked by printing the computation graph and increasing the nb_updates should increase the computation graph but it doesn’t. I believe it has something to do with the fact we are changing/mutating the OrderDict/params and looping through it at the same time. I tried to fix it by collecting the parameters before and after running the loop but the somehow it thinks its empty for the second loop (i.e. t=1 doesn’t print anything when it should):
import torch
import torch.nn as nn

from torchviz import make_dot

import copy

from collections import OrderedDict

def del_attr(obj, names):
    if len(names) == 1:
        delattr(obj, names[0])
    else:
        del_attr(getattr(obj, names[0]), names[1:])
        
def set_attr(obj, names, val):
    if len(names) == 1:
        setattr(obj, names[0], val)
    else:
        set_attr(getattr(obj, names[0]), names[1:], val)

        
x = torch.randn(1, requires_grad=True)
y = torch.randn(1, requires_grad=True)

criterion = nn.CrossEntropyLoss()

loss_net = nn.Sequential(OrderedDict([('l_fc0', nn.Linear(in_features=1,out_features=1, bias=True))]))
loss_net.l_fc0.weight.requires_grad=False
loss_net.l_fc0.bias.requires_grad=False

hidden = torch.randn(size=(1,1),requires_grad=True)
updater_net = nn.Sequential(OrderedDict([('u_fc0',nn.Linear(in_features=1,out_features=1))]))
updater_net.u_fc0.bias.requires_grad = False
print(f'updater_net.u_fc0.weight.is_leaf = {updater_net.u_fc0.weight.is_leaf}')
#
outputs_virgin = loss_net(x)
params = dict(dict(loss_net.named_parameters()),**{'x':x})
make_dot(outputs_virgin, params=params).render('loss_net_x', format='png')
#
nb_updates = 2
params = list(loss_net.named_parameters())
for t in range(nb_updates):
    print(f't = {t}')
    ## w^<t> := f(w^<t-1>,delta^<t-1>)
    for (name, w) in params:
        delta = updater_net(hidden).view(1)
        wt = w + delta
        print(f'w^<{t}> = {wt}')
        del_attr(loss_net, name.split("."))
        set_attr(loss_net, name.split("."), wt)
    params = list(loss_net.named_parameters())

print()
print(f'updater_net.u_fc0.weight.is_leaf = {updater_net.u_fc0.weight.is_leaf}')
print(f'loss_net.l_fc0.weight.is_leaf = {loss_net.l_fc0.weight.is_leaf}')

outputs = loss_net(x)
loss_val = (outputs - y)**2
loss_val.backward()
print()
print(f'-- params that dont matter if they have gradients --')
print(f'loss_net.grad = {loss_net.l_fc0.weight.grad}')
print('-- params we want to have gradients --')
print(f'hidden.grad = {hidden.grad}') # None because this is not a leaf, it is overriden in the for loop above.
print(f'updater_net.u_fc0.weight.grad = {updater_net.u_fc0.weight.grad}')
#print(f'updater_net.u_fc0.bias.grad = {updater_net.u_fc0.bias.grad}')
params = dict(dict(updater_net.named_parameters()),**{'x':x,'y':y,'hidden':hidden},**dict(loss_net.named_parameters()))
make_dot(loss_val, params=params).render('loss_val', format='png')