def lamb_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 = []
        grads = []
        exp_avgs = []
        exp_avg_sqs = []
        trails = []
        state_steps = []

        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)
                if grad.is_sparse:
                    raise RuntimeError("Lamb does not support sparse gradients")
                grads.append(grad)

                state = self.state[p]
                # Lazy state initialization
                if len(state) == 0:
                    state["step"] = 0
                    buffer_dtype = p.dtype if p.dtype is torch.float64 else torch.float
                    state["exp_avg"] = torch.zeros(
                        p.shape, dtype=buffer_dtype, device=p.device
                    )
                    state["exp_avg_sq"] = torch.zeros(
                        p.shape, dtype=buffer_dtype, device=p.device
                    )

                exp_avgs.append(state["exp_avg"])
                exp_avg_sqs.append(state["exp_avg_sq"])

                # update the steps for each param group update
                state["step"] += 1
                # record the step after step update
                state_steps.append(state["step"])

        beta1, beta2 = group["betas"]
        _lamb_fused_impl(
            params_with_grad,
            grads,
            exp_avgs,
            exp_avg_sqs,
            self.params_attr,
            state_steps,
            beta1,
            beta2,
            group["lr"],
            group["weight_decay"],
            group["eps"],
        )
    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