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