The code from my last post is wrong (doesn’t work as mentioned). Keeping a list of optimizers is useless, we should keep a (actually two) list of models and copy the gradients. The following is the correct code.
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class DNN(nn.Module):
    def __init__(self,dim):
        super(DNN,self).__init__()
        self.fc=nn.Sequential(
            nn.Linear(dim,dim),
            nn.Tanh(),
            nn.Linear(dim,dim)
        )
    def forward(self,data):
        return self.fc(data)

def main():
    E=50
    T=20
    D=5
    dim=4
    model=DNN(dim)
    optimizer=torch.optim.Adam(model.parameters())
    loss_fn=nn.MSELoss()
    data_in=[[torch.zeros(dim) for t in range(T)] for e in range(E)]
    data_out=[[torch.zeros(dim) for t in range(T)] for e in range(E)]
    data_target=[[torch.zeros(dim) for t in range(T)] for e in range(E)]
    model_list=[None for e in range(E)]
    model_update_list=[None for e in range(E)]
    for e in range(E):
        if e==0:
            model_list[e]=DNN(dim)
        else:
            model_list[e]=type(model_update_list[e-1])(dim)
            model_list[e].load_state_dict(model_update_list[e-1].state_dict())
        for t in range(T):
            data_in[e][t]=torch.from_numpy(np.random.rand(dim)).float()
            data_out[e][t]=model_list[e](data_in[e][t])
            data_target[e][t]=torch.from_numpy(np.random.rand(dim)).float()
        loss=0.0
        for t in range(T):
            loss=loss_fn(data_target[e][t],data_out[e][t])
        optimizer=torch.optim.Adam(model_list[e].parameters())
        optimizer.zero_grad()
        loss.backward()
        model_update_list[e]=type(model_list[e])(dim)
        model_update_list[e].load_state_dict(model_list[e].state_dict())
        optimizer_update=torch.optim.Adam(model_update_list[e].parameters())
        if e>=D:
            for p1,p2 in zip(model_update_list[e].named_parameters(),model_list[e-D].named_parameters()):
                p1[1].grad=p2[1].grad
            optimizer_update.step()
    
if __name__=='__main__':
    main()