def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('Adamax does not support sparse gradients')

                p_data_fp32 = p.data.float()

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p_data_fp32)
                    state['exp_inf'] = torch.zeros_like(p_data_fp32)
                else:
                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
                    state['exp_inf'] = state['exp_inf'].type_as(p_data_fp32)

                exp_avg, exp_inf = state['exp_avg'], state['exp_inf']
                beta1, beta2 = group['betas']
                eps = group['eps']

                state['step'] += 1

                # Update biased first moment estimate.
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                # Update the exponentially weighted infinity norm.
                torch.max(
                    exp_inf.mul_(beta2),
                    grad.abs_(),
                    out=exp_inf,
                )

                step_size = group['lr']
                if group['bias_correction']:
                    bias_correction = 1 - beta1 ** state['step']
                    step_size /= bias_correction

                if group['weight_decay'] != 0:
                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)

                p_data_fp32.addcdiv_(-step_size, exp_avg, exp_inf.add(eps))

                p.data.copy_(p_data_fp32)

        return loss
Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        
infoxlm/fairseq/fairseq/optim/adamax.py