However, I create a simple toy linear regression code; it seems to me it doesn’t require retain_graph=True. However, in my true code; which are very similar; requires retain_graph=True
This is quite confusing that I’m not sure if I fully understand how does backward() works
x = torch.randn(100,1)
y = 1 + 2*x + torch.randn(x.shape[0]).reshape(100,1)
loss = lambda x: ((x-y).pow(2)).mean()
n_dim = 1
F = torch.nn.Sequential()



for l in range(0, 10):
    F.add_module("%d"%l, nn.Linear(n_dim, 1))
    
    if l == 0: # first layer create optimizer
        optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, F.parameters()), lr=1e-3)
    else: # new layers added to optimizer
        params = [param for name, param in F.named_parameters() if "%d." % l in name]
        optimizer.add_param_group({"params": params})
        
        
    # freeze
    for name, param in F.named_parameters():
        if "%d." % l not in name:
            param.requires_grad = False

    # train new layers only
    for _ in tqdm(range(100)):

        optimizer.zero_grad()
        pred_y = F(x)
        #loss = ((pred_y-y).pow(2)).mean()
        loss = criteria(pred_y)
        loss.backward()

        optimizer.zero_grad()
        
        
    # unfreeze
    for param in F.parameters():
        param.requires_grad = True

    # train on all layers
    for _ in tqdm(range(100)):
        optimizer.zero_grad()
        pred_y = F(x)
        #loss = ((pred_y-y).pow(2)).mean()
        loss = criteria(pred_y)
        loss.backward()

        optimizer.zero_grad()