def lars_step(self, closure=None):
    """Performs a single optimization step.
    Args:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()

    for group in self.param_groups:
        params_with_grad = []
        params2 = []
        d_p_list = []
        momentum_buffer_list = []
        has_sparse_grad = False

        for p in group["params"]:
            grad = (
                get_bf16_grad(p, self.params_attr)
                if is_master_weight(p, self.params_attr)
                else p.grad
            )
            if grad is not None:
                params_with_grad.append(p)
                d_p_list.append(grad)
                if grad.is_sparse:
                    has_sparse_grad = True

                state = self.state[p]
                if "momentum_buffer" not in state:
                    momentum_buffer_list.append(None)
                else:
                    momentum_buffer_list.append(state["momentum_buffer"])

                param2 = get_param2(p, self.params_attr)
                params2.append(param2)
        if group["lars"]:
            lars(
                params_with_grad,
                params2,
                d_p_list,
                momentum_buffer_list,
                eeta=group["eeta"],
                eps=group["epsilon"],
                weight_decay=group["weight_decay"],
                momentum=group["momentum"],
                lr=group["lr"],
                dampening=0,
                nesterov=0,
                maximize=0,
                has_sparse_grad=has_sparse_grad,
                foreach=None,
                fused=self.fused,
            )
        else:
            sgd(
                params_with_grad,
                params2,
                d_p_list,
                momentum_buffer_list,
                weight_decay=group["weight_decay"],
                momentum=group["momentum"],
                lr=group["lr"],
                dampening=0,
                nesterov=0,
                maximize=0,
                has_sparse_grad=has_sparse_grad,
                foreach=None,
                fused=self.fused,
            )

        # update momentum_buffers in state
        for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list):
            state = self.state[p]
            state["momentum_buffer"] = momentum_buffer

    return loss
Performs a single optimization step.
    Args:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    
intel_extension_for_pytorch/optim/_functional.py