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
                self.bf16 = p.grad.dtype == torch.bfloat16
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError(
                        "Lamb does not support sparse gradients, consider SparseAdam instad."
                    )

                state = self.state[p]
                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state["exp_avg_sq"] = torch.zeros_like(p.data)
                    # Lower bits for bf16 params
                    if p.data.dtype == torch.bfloat16:
                        state["low_bits"] = torch.zeros_like(p.data)
                    state["weight_norm"] = -1.0

                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                if p.data.dtype == torch.bfloat16:
                    low_bits = state["low_bits"]
                beta1, beta2 = group["betas"]
                weight_norm = state["weight_norm"]

                state["step"] += 1

                if p.data.dtype == torch.bfloat16:
                    AssertionError(False, "BF16 LAMB optimizer not implemented yet!")
                    state["weight_norm"] = ipex_cpp.tpp_fused_split_lamb(
                        p.data,
                        low_bits,
                        grad.contiguous(),
                        exp_avg,
                        exp_avg_sq,
                        beta1,
                        beta2,
                        weight_norm,
                        group["lr"],
                        group["weight_decay"],
                        group["eps"],
                    )
                else:
                    state["weight_norm"] = ipex_cpp.tpp_fused_lamb(
                        p.data,
                        grad.contiguous(),
                        exp_avg,
                        exp_avg_sq,
                        beta1,
                        beta2,
                        weight_norm,
                        group["lr"],
                        group["weight_decay"],
                        group["eps"],
                    )

                #  # Decay the first and second moment running average coefficient
                #  if self.bf16:
                #      # m_t
                #      exp_avg.mul_(beta1).add_(1 - beta1, grad_fp32)
                #      # v_t
                #      exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad_fp32, grad_fp32)
                #  else:
                #      # m_t
                #      exp_avg.mul_(beta1).add_(1 - beta1, grad)
                #      # v_t
                #      exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                #  # Paper v3 does not use debiasing.
                #  # bias_correction1 = 1 - beta1 ** state['step']
                #  # bias_correction2 = 1 - beta2 ** state['step']
                #  # Apply bias to lr to avoid broadcast.
                #  step_size = group['lr'] # * math.sqrt(bias_correction2) / bias_correction1

                #  weight_norm = data_fp32.pow(2).sum().sqrt().clamp(0, 10) if self.bf16 \
                #      else p.data.pow(2).sum().sqrt().clamp(0, 10)

                #  adam_step = exp_avg / exp_avg_sq.sqrt().add(group['eps'])
                #  if group['weight_decay'] != 0:
                #      if self.bf16:
                #          adam_step.add_(group['weight_decay'], data_fp32)
                #      else:
                #          adam_step.add_(group['weight_decay'], p.data)

                #  adam_norm = adam_step.pow(2).sum().sqrt()
                #  if weight_norm == 0 or adam_norm == 0:
                #      trust_ratio = 1
                #  else:
                #      trust_ratio = weight_norm / adam_norm
                #  state['weight_norm'] = weight_norm
                #  state['adam_norm'] = adam_norm
                #  state['trust_ratio'] = trust_ratio
                #  if self.adam:
                #      trust_ratio = 1

                #  if self.bf16:
                #      data_fp32.add_(-step_size * trust_ratio, adam_step)
                #      p.data = data_fp32.to(torch.bfloat16)
                #  else:
                #      p.data.add_(-step_size * trust_ratio, adam_step)

        return loss
Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        
intel_extension_for_pytorch/cpu/tpp/optim.py