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.
docstr: Performs a single optimization step.
  
  Arguments:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: 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:
          with torch.enable_grad():
              loss = closure()
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              grad = p.grad
              if grad.is_sparse:
                  raise RuntimeError(
                      "LAMB does not support sparse gradients"
                  )
              amsgrad = group["amsgrad"]
              exclude_from_layer_adaptation = group[
                  "exclude_from_layer_adaptation"
              ]
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(
                      p, memory_format=torch.preserve_format
                  )
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(
                      p, memory_format=torch.preserve_format
                  )
                  if amsgrad:
                      # Maintains max of all exp. moving avg. of sq. grad. values
                      state["max_exp_avg_sq"] = torch.zeros_like(
                          p, memory_format=torch.preserve_format
                      )
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
              if amsgrad:
                  max_exp_avg_sq = state["max_exp_avg_sq"]
              beta1, beta2 = group["betas"]
  
              state["step"] += 1
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
  
              # Decay the first and second moment running average coefficient
              exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
              exp_avg_sq.mul_(beta2).addcmul_(
                  grad, grad, value=1 - beta2
              )
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch.max(
                      max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq
                  )
                  # Use the max. for normalizing running avg. of gradient
                  denom = (
                      max_exp_avg_sq.sqrt()
                      / math.sqrt(bias_correction2)
                  ).add_(group["eps"])
              else:
                  denom = (
                      exp_avg_sq.sqrt() / math.sqrt(bias_correction2)
                  ).add_(group["eps"])
  
              numerator = exp_avg / bias_correction1
              update = numerator / denom
  
              if group["weight_decay"] != 0:
                  update = update.add(
                      p.data, alpha=group["weight_decay"]
                  )
  
              trust_ratio = 1.0
              if not exclude_from_layer_adaptation:
                  w_norm = torch.norm(p.data)
                  g_norm = torch.norm(update)
  
                  if w_norm > 0 and g_norm > 0:
                      trust_ratio = w_norm / g_norm
  
              p.add_(update, alpha=-group["lr"] * trust_ratio)
  
      return loss
  
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 3 for calling the LAMB obj's underlying member method code with example usage: obj.step(closure) and returns: loss