def _lamb_fused_impl(
    params: List[Tensor],
    grads: List[Tensor],
    exp_avgs: List[Tensor],
    exp_avg_sqs: List[Tensor],
    attr: dict,
    state_steps: List[int],
    beta1: float,
    beta2: float,
    lr: float,
    weight_decay: float,
    eps: float,
):
    r"""Functional API that performs Lamb algorithm computation.
    See :class:`~torch.optim.Lamb` for details.
    """

    for i, param in enumerate(params):
        grad = grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_sq = exp_avg_sqs[i]
        step = state_steps[i]
        param2 = get_param2(param, attr)
        torch.ops.torch_ipex.lamb_fused_step(
            param,
            exp_avg,
            exp_avg_sq,
            grad,
            param2,
            step,
            beta1,
            beta2,
            lr,
            weight_decay,
            eps,
        )
Functional API that performs Lamb algorithm computation.
    See :class:`~torch.optim.Lamb` for details.
    
intel_extension_for_pytorch/optim/_functional.py