code: def construct_optimizer(
      model: torch.nn.Module, cfg: OptimizerConf  # noqa
  ) -> torch.optim.Optimizer:
      """
      Constructs a stochastic gradient descent or ADAM (or ADAMw) optimizer
      with momentum. i.e, constructs a torch.optim.Optimizer with zero-weight decay
      Batchnorm and/or no-update 1-D parameters support, based on the config.
  
      Supports wrapping the optimizer with Layer-wise Adaptive Rate Scaling
      (LARS): https://arxiv.org/abs/1708.03888
  
      Args:
          model (nn.Module): model to perform stochastic gradient descent
              optimization or ADAM optimization.
          cfg (OptimizerConf): Hydra/Omega conf object consisting hyper-parameters
              of SGD or ADAM, includes base learning rate,  momentum, weight_decay,
              dampening and etc. The supported config schema is `OptimizerConf`.
              Example config files can be found at,
              `pytorchvideo_trainer/conf/module/optim`
      """
      bn_parameters = []
      non_bn_parameters = []
      zero_parameters = []
      no_grad_parameters = []
      skip = {}
  
      if hasattr(model, "no_weight_decay"):
          skip = model.no_weight_decay()  # pyre-ignore[29]
  
      for name, m in model.named_modules():
          is_bn = isinstance(m, torch.nn.modules.batchnorm._NormBase)
          for p in m.parameters(recurse=False):
              if not p.requires_grad:
                  no_grad_parameters.append(p)
              elif is_bn:
                  bn_parameters.append(p)
              elif name in skip:
                  zero_parameters.append(p)
              elif cfg.zero_weight_decay_1d_param and (
                  len(p.shape) == 1 or name.endswith(".bias")
              ):
                  zero_parameters.append(p)
              else:
                  non_bn_parameters.append(p)
  
      optim_params = [
          {
              "params": bn_parameters,
              "weight_decay": cfg.bn_weight_decay,
              "apply_LARS": False,
          },
          {
              "params": non_bn_parameters,
              "weight_decay": cfg.weight_decay,
              "apply_LARS": cfg.lars_on,
          },
          {
              "params": zero_parameters,
              "weight_decay": 0.0,
              "apply_LARS": cfg.lars_on,
          },
      ]
      optim_params = [
          x for x in optim_params if len(x["params"])
      ]  # pyre-ignore[6]
  
      # Check all parameters will be passed into optimizer.
      assert len(list(model.parameters())) == len(
          non_bn_parameters
      ) + len(bn_parameters) + len(zero_parameters) + len(
          no_grad_parameters
      ), "parameter size does not match: {} + {} + {} + {} != {}".format(
          len(non_bn_parameters),
          len(bn_parameters),
          len(zero_parameters),
          len(no_grad_parameters),
          len(list(model.parameters())),
      )
      print(
          "bn {}, non bn {}, zero {} no grad {}".format(
              len(bn_parameters),
              len(non_bn_parameters),
              len(zero_parameters),
              len(no_grad_parameters),
          )
      )
  
      if cfg.method == "sgd":
          optimizer = torch.optim.SGD(
              optim_params,
              lr=cfg.lr,
              momentum=cfg.momentum,
              weight_decay=cfg.weight_decay,
              dampening=cfg.dampening,
              nesterov=cfg.nesterov,
          )
      elif cfg.method == "adam":
          optimizer = torch.optim.Adam(
              optim_params,
              lr=cfg.lr,
              betas=(0.9, 0.999),
              weight_decay=cfg.weight_decay,
          )
      elif cfg.method == "adamw":
          optimizer = torch.optim.AdamW(
              optim_params,
              lr=cfg.lr,
              eps=1e-08,
              weight_decay=cfg.weight_decay,
          )
      else:
          raise NotImplementedError(
              "Does not support {} optimizer".format(cfg.method)
          )
  
      if cfg.lars_on:
          optimizer = LARS(
              optimizer=optimizer, trust_coefficient=0.001, clip=False
          )
      return optimizer
  
imports: from dataclasses import dataclass
  import torch
  from hydra.core.config_store import ConfigStore
  from omegaconf import MISSING
label_desc: Use image node_id 1 for calling a global function with example usage: construct_optimizer(model, cfg) and returns: optimizer