Here is the main part of my code
    def __init__(self, T: int, dim_size: int,lag_list):
        super(MF, self).__init__()
        self.t_embeddings = nn.Embedding(T, dim_size）
        
        self.modlst=torch.nn.ModuleList([])
        self.lag_list=lag_list
        self.l=len(lag_list)
        for i in range(self.l):
             self.modlst.append(torch.nn.Linear(dim_size,dim_size,bias=None))

    def forward(self, times):
         mt = self.t_embeddings(times)
         return(mt)

     def var(self, times):
        tmp = self.t_embeddings(times)
                
        for i in range(max(self.lag_list),tmp.shape[0]):
            for j in range(self.l):
                tmp[i]=tmp[i]+self.modlst[j](tmp[i-self.lag_list[j]])

        return(tmp)
            
def ARLoss(model,T,lag_list):
    times=torch.arange(T).to(device)
    ar_y=model.var(times)
    y=model(times,lag_list)
    loss=torch.linalg.norm(ar_y-y)
    return loss

model=MF(t,lag_list)
loss_func=torch.nn.MSEloss()
loss=loss_func(target,model(times,lag_list))+ARLoss(model,T,lag_list)

......