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()

        self._one_time_setup()
        self._step += 1
        self.state[self.param_groups[0]["params"][0]]["step"] = self._step
        self.merge_acc_grad()
        if self.perform_allreduce:
            self.sync_grads()

        for ii, fp in enumerate(self.flat_params):
            group = fp.group
            beta1, beta2 = group["betas"]
            lr = group["lr"]
            eps = group["eps"]
            weight_decay = group["weight_decay"]
            ipex_cpp.tpp_fused_lamb_v2(
                fp._flat_w,
                fp._flat_g,
                fp._flat_m,
                fp._flat_v,
                fp._flat_u,
                fp._flat_wl,
                fp._offsets,
                fp._block_sizes,
                fp._block2param,
                fp._weight_norms,
                fp._update_norms,
                weight_decay,
                beta1,
                beta2,
                lr,
                eps,
                fp.block_size,
                self._step,
                self.fused_param_norm,
            )
            # if weight_decay > 0.0 and torch.distributed.get_rank() < 2:
            #   print(f"wn: {fp._weight_norms[:5].sqrt()}  un: {fp._update_norms[:5].sqrt()}")
            # if weight_decay > 0.0:
            #   print(f"XXX {self._step:3d} NORM {ii}: wn: \
            #   {fp._weight_norms[0].sqrt().item():.10f}  un: {fp._update_norms[0].sqrt().item():.10f}")

        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