def METHOD_NAME(self):
    for group in self.param_groups:
        group['step'] = 0
        for p in group['params']:
            if p.requires_grad:
                state = self.state[p]
                # State initialization
                # Exponential moving average of gradient values
                state['exp_avg'] = torch.zeros_like(p)
                # Exponential moving average of squared gradient values
                state['exp_avg_sq'] = torch.zeros_like(p)
                # Exponential moving average of gradient difference
                state['exp_avg_diff'] = torch.zeros_like(p)