code: def _load_optimizer(model, cfg_train, cfg_impl, initial_time):
      # Filter some parameters
      grouped_parameters = group_parameters(model, cfg_train)
  
      # Select optimizer implementation
      if cfg_train.optim.type == "AdamW":
          optimizer_class = torch.optim.AdamW
      elif cfg_train.optim.type == "Adam":
          optimizer_class = torch.optim.Adam
      elif cfg_train.optim.type == "RAdam":
          optimizer_class = torch.optim.RAdam
      elif cfg_train.optim.type == "SGD":
          optimizer_class = torch.optim.SGD
      elif cfg_train.optim.type == "Adafactor":
          optimizer_class = transformers.Adafactor
      elif cfg_train.optim.type == "Shampoo":
          optimizer_class = Shampoo
      elif cfg_train.optim.type == "AdaHessian":
          optimizer_class = Adahessian
      elif cfg_train.optim.type == "AdamWScale":
          optimizer_class = AdamWScale
      elif cfg_train.optim.type == "Sophia-G":
          optimizer_class = Sophia
      elif cfg_train.optim.type == "Lion":
          from lion_pytorch import Lion
  
          optimizer_class = Lion
  
      elif cfg_train.optim.type == "Adam8bit":
          import bitsandbytes as bnb
  
          optimizer_class = bnb.optim.Adam8bit
      elif cfg_train.optim.type == "AGD":
          depth = len(list(model.parameters()))
          optimizer_class = partial(AGD, depth=depth)
      else:
          raise ValueError(
              f"Invalid optimizer {cfg_train.optim.type} given."
          )
      optimizer_args = {
          k: v for k, v in cfg_train.optim.items() if k != "type"
      }
      if (
          cfg_impl.foreach_optimizer
          and cfg_train.optim.type != "Shampoo"
      ):
          optimizer_args["foreach"] = True
  
      if (
          torch.distributed.is_initialized()
          and cfg_impl.zero_redundancy_optimizer
      ):
          # The overlap option is a whole bucket of problems in itself for now...
          optimizer = ZeroRedundancyOptimizer(
              grouped_parameters,
              optimizer_class=optimizer_class,
              parameters_as_bucket_view=True,
              overlap_with_ddp=False,
              **optimizer_args,
          )
      else:
          optimizer = optimizer_class(
              grouped_parameters, **optimizer_args
          )
  
      if cfg_train.optim_mod.name == "none":
          optimizer_to_schedule = optimizer
      else:
          optim_params = {
              k: v
              for k, v in cfg_train.optim_mod.items()
              if k != "name"
          }
          if cfg_train.optim_mod.name == "LARS":
              optimizer = LARS(optimizer, **optim_params)
          elif cfg_train.optim_mod.name == "LARC":
              optimizer = LARS(optimizer, **optim_params)
          elif cfg_train.optim_mod.name == "SAM":
              optimizer = SAM(optimizer, **optim_params)
          elif cfg_train.optim_mod.name == "progressive-batching":
              optimizer = ProgressiveBatching(optimizer, **optim_params)
  
          optimizer_to_schedule = optimizer.optim
  
      scheduler = get_schedule_fn(initial_time, cfg_train)(
          optimizer_to_schedule
      )
  
      return optimizer, scheduler
  
imports: import torch
  import torch._inductor.utils
  import os
  import json
  from omegaconf import OmegaConf
  from functools import partial
  from contextlib import nullcontext
  import time
  import logging
  import transformers
  from safetensors.torch import load_file, save_file
  from transformers.utils.generic import working_or_temp_dir
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from .utils import group_parameters, prepare_pretraining_dataloader, update_ema, updated_latest_weight_average
  from .optimizers.schedulers import get_schedule_fn
  from .optimizers import Adahessian, AdamWScale, Shampoo, LARS, SAM, ProgressiveBatching, AGD, Sophia
  import warnings
label_desc: Use image node_id 2 for calling a global function with example usage: _load_optimizer(model, cfg_train, cfg_impl, initial_time) and returns: optimizer,  scheduler