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:
            weight_decay = group["weight_decay"]
            momentum = group["momentum"]
            dampening = group["dampening"]
            nesterov = group["nesterov"]

            for p in group["params"]:
                if p.grad is None:
                    continue
                d_p = p.grad.data
                if p.dtype == torch.bfloat16:
                    param_state = self.state[p]
                    if "low_bits" not in param_state:
                        buf = param_state["low_bits"] = torch.zeros_like(
                            p.data, dtype=torch.short
                        )
                    else:
                        buf = param_state["low_bits"]
                # if weight_decay != 0:
                #     d_p = d_p.add(weight_decay, p.data)
                # if momentum != 0:
                #     param_state = self.state[p]
                #     if 'momentum_buffer' not in param_state:
                #         buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
                #     else:
                #         buf = param_state['momentum_buffer']
                #         buf.mul_(momentum).add_(1 - dampening, d_p)
                #     if nesterov:
                #         d_p = d_p.add(momentum, buf)
                #     else:
                #         d_p = buf

                # p.data.add_(-group['lr'], d_p)
                if p.dtype == torch.bfloat16:
                    ipex_cpp.tpp_bf16_split_add_(p.data, buf, d_p, -group["lr"])
                else:
                    if d_p.is_sparse:
                        ipex_cpp.tpp_dense_sparse_add(p.data, d_p, -group["lr"])
                    else:
                        p.data.add_(d_p, alpha=-group["lr"])

        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