def restore_grad(self):
    for (name, param) in self.model.named_parameters():
        if param.requires_grad:
            param.grad = self.grad_backup[name]
def restore_grad(self):
    for (name, param) in self.model.named_parameters():
        if param.requires_grad:
            param.grad = self.grad_backup[name]