how to self-define the autograd to realize the CNN + SVRG

Hi , all.
I want modify the way of grandient update.
After the CNN model is setted, when it comes to training, we need to update the gradient by the following algorithm in picture

So, I don’t know how to calculate the average gradient of all samples under the last period’s parameter, and how to transfer the gradient (the fomula w_t in the Line 9 of the picture ) to the model

the follow is my code and it remains to be modified.

# 主要的训练过程
for epoch in range(n_epoch) :
    for data in img_dataloader :
        img = data
        img = img.cpu()  # 读入到gpu
        output1, output = model(img)  # 编码结果、解码结果
        loss = criterion(output, img)  # MSE
        optimizer.zero_grad()  # 梯度置零
        loss.backward()  # 梯度回传
        #after calculate the ori grad and before update the next parameter, we need to amend the gradient
        #get current grad
        #amended grad = current_grad - mu_grad
        #but I dont know how to transfer the amended grad to model
        optimizer.step()  # 更新参数
        #there maybe need to save the last period's gradient but I dont know how to do
        #parameter = last paramater
        #mu_grad = calcultegrad(sample,last_parameter).mean()