def step(self, closure=None):
    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('Ralamb does not support sparse gradients')
            p_data_fp32 = p.data.float()
            state = self.state[p]
            if (len(state) == 0):
                state['step'] = 0
                state['exp_avg'] = torch.zeros_like(p_data_fp32)
                state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
            else:
                state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
                state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
            (exp_avg, exp_avg_sq) = (state['exp_avg'], state['exp_avg_sq'])
            (beta1, beta2) = group['betas']
            exp_avg.mul_(beta1).add_((1 - beta1), grad)
            exp_avg_sq.mul_(beta2).addcmul_((1 - beta2), grad, grad)
            state['step'] += 1
            buffered = self.buffer[int((state['step'] % 10))]
            if (state['step'] == buffered[0]):
                (N_sma, radam_step_size) = (buffered[1], buffered[2])
            else:
                buffered[0] = state['step']
                beta2_t = (beta2 ** state['step'])
                N_sma_max = ((2 / (1 - beta2)) - 1)
                N_sma = (N_sma_max - (((2 * state['step']) * beta2_t) / (1 - beta2_t)))
                buffered[1] = N_sma
                if (N_sma >= 5):
                    radam_step_size = ((group['lr'] * math.sqrt((((((((1 - beta2_t) * (N_sma - 4)) / (N_sma_max - 4)) * (N_sma - 2)) / N_sma) * N_sma_max) / (N_sma_max - 2)))) / (1 - (beta1 ** state['step'])))
                else:
                    radam_step_size = (group['lr'] / (1 - (beta1 ** state['step'])))
                buffered[2] = radam_step_size
            if (group['weight_decay'] != 0):
                p_data_fp32.add_(((- group['weight_decay']) * group['lr']), p_data_fp32)
            radam_step = p_data_fp32.clone()
            if (N_sma >= 5):
                denom = exp_avg_sq.sqrt().add_(group['eps'])
                radam_step.addcdiv_((- radam_step_size), exp_avg, denom)
            else:
                radam_step.add_((- radam_step_size), exp_avg)
            radam_norm = radam_step.pow(2).sum().sqrt()
            weight_norm = p.data.pow(2).sum().sqrt().clamp(0, 10)
            if (0 in (weight_norm, radam_norm)):
                trust_ratio = 1
            else:
                trust_ratio = (weight_norm / radam_norm)
            state['weight_norm'] = weight_norm
            state['adam_norm'] = radam_norm
            state['trust_ratio'] = trust_ratio
            if (N_sma >= 5):
                p_data_fp32.addcdiv_(((- radam_step_size) * trust_ratio), exp_avg, denom)
            else:
                p_data_fp32.add_(((- radam_step_size) * trust_ratio), exp_avg)
            p.data.copy_(p_data_fp32)
    return loss
def step(self, closure=None):
    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('Ralamb does not support sparse gradients')
            p_data_fp32 = p.data.float()
            state = self.state[p]
            if (len(state) == 0):
                state['step'] = 0
                state['exp_avg'] = torch.zeros_like(p_data_fp32)
                state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
            else:
                state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
                state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
            (exp_avg, exp_avg_sq) = (state['exp_avg'], state['exp_avg_sq'])
            (beta1, beta2) = group['betas']
            exp_avg.mul_(beta1).add_((1 - beta1), grad)
            exp_avg_sq.mul_(beta2).addcmul_((1 - beta2), grad, grad)
            state['step'] += 1
            buffered = self.buffer[int((state['step'] % 10))]
            if (state['step'] == buffered[0]):
                (N_sma, radam_step_size) = (buffered[1], buffered[2])
            else:
                buffered[0] = state['step']
                beta2_t = (beta2 ** state['step'])
                N_sma_max = ((2 / (1 - beta2)) - 1)
                N_sma = (N_sma_max - (((2 * state['step']) * beta2_t) / (1 - beta2_t)))
                buffered[1] = N_sma
                if (N_sma >= 5):
                    radam_step_size = ((group['lr'] * math.sqrt((((((((1 - beta2_t) * (N_sma - 4)) / (N_sma_max - 4)) * (N_sma - 2)) / N_sma) * N_sma_max) / (N_sma_max - 2)))) / (1 - (beta1 ** state['step'])))
                else:
                    radam_step_size = (group['lr'] / (1 - (beta1 ** state['step'])))
                buffered[2] = radam_step_size
            if (group['weight_decay'] != 0):
                p_data_fp32.add_(((- group['weight_decay']) * group['lr']), p_data_fp32)
            radam_step = p_data_fp32.clone()
            if (N_sma >= 5):
                denom = exp_avg_sq.sqrt().add_(group['eps'])
                radam_step.addcdiv_((- radam_step_size), exp_avg, denom)
            else:
                radam_step.add_((- radam_step_size), exp_avg)
            radam_norm = radam_step.pow(2).sum().sqrt()
            weight_norm = p.data.pow(2).sum().sqrt().clamp(0, 10)
            if (0 in (weight_norm, radam_norm)):
                trust_ratio = 1
            else:
                trust_ratio = (weight_norm / radam_norm)
            state['weight_norm'] = weight_norm
            state['adam_norm'] = radam_norm
            state['trust_ratio'] = trust_ratio
            if (N_sma >= 5):
                p_data_fp32.addcdiv_(((- radam_step_size) * trust_ratio), exp_avg, denom)
            else:
                p_data_fp32.add_(((- radam_step_size) * trust_ratio), exp_avg)
            p.data.copy_(p_data_fp32)
    return loss