class_docstr: Extends ADAM in pytorch to incorporate LAMB algorithm from the paper: `Large batch optimization for deep
  learning: Training BERT in 76 minutes <https://arxiv.org/pdf/1904.00962.pdf>`_.
  
  Args:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr (float): learning rate
      betas (Tuple[float, float], optional): coefficients used for computing
          running averages of gradient and its square (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve numerical stability (default: 1e-8)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
      exclude_from_layer_adaptation (bool, optional): layers which do not need LAMB
          layer adaptation (default: False)
      amsgrad (boolean, optional): whether to use the AMSGrad variant of this
          algorithm from the paper `On the Convergence of Adam and Beyond <https://arxiv.org/pdf/1904.09237.pdf>`_
          (default: False)
  
  Example:
      >>> from torch import nn
      >>> model = nn.Linear(10, 1)
      >>> optimizer = LAMB(model.parameters(), lr=0.1)
      >>> optimizer.zero_grad()
      >>> # loss_fn(model(input), target).backward()
      >>> optimizer.step()
  
  .. warning::
      Since the default weight decay for LAMB is set to 0., we do not club together
      0. weight decay and exclusion from layer adaptation like LARS. This would cause
      the optimizer to exclude all layers from layer adaptation.
code: def __init__(
      self,
      params,
      lr: float = 1e-3,
      betas: Tuple[float, float] = (0.9, 0.999),
      eps: float = 1e-6,
      weight_decay: float = 0,
      exclude_from_layer_adaptation: bool = False,
      amsgrad: bool = False,
  ):
      if not lr >= 0.0:
          raise ValueError(f"Invalid learning rate: {lr}")
      if not eps >= 0.0:
          raise ValueError(f"Invalid epsilon value: {eps}")
      if not 0.0 <= betas[0] < 1.0:
          raise ValueError(
              f"Invalid beta parameter at index 0: {betas[0]}"
          )
      if not 0.0 <= betas[1] < 1.0:
          raise ValueError(
              f"Invalid beta parameter at index 1: {betas[1]}"
          )
      if not weight_decay >= 0.0:
          raise ValueError(
              f"Invalid weight_decay value: {weight_decay}"
          )
      defaults = {
          "lr": lr,
          "betas": betas,
          "eps": eps,
          "weight_decay": weight_decay,
          "exclude_from_layer_adaptation": exclude_from_layer_adaptation,
          "amsgrad": amsgrad,
      }
      super().__init__(params, defaults)
  
file_path: lightning-flash/src/flash/core/optimizers/lamb.py
imports: import math
  from typing import Tuple
  import torch
  from torch.optim.optimizer import Optimizer
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE
label_desc: Use image node_id 1 to create a new LAMB object from inherited base classes: Optimizer with example: obj = LAMB(params, lr, betas, eps, weight_decay, exclude_from_layer_adaptation, amsgrad)