code: def get_optimizer_param_scheduler(optimizer):
      """Build the learning rate scheduler."""
      args = get_args()
  
      # Iteration-based training.
      if args.train_iters:
          if args.lr_decay_iters is None:
              args.lr_decay_iters = args.train_iters
          lr_decay_steps = args.lr_decay_iters * args.global_batch_size
          wd_incr_steps = args.train_iters * args.global_batch_size
          if args.lr_warmup_fraction is not None:
              lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
          else:
              lr_warmup_steps = (
                  args.lr_warmup_iters * args.global_batch_size
              )
      # Sample-based training.
      elif args.train_samples:
          # We need to set training iters for later use. Technically
          # we need to adjust the training samples too (due to last
          # batch being incomplete) but we leave it as is for now.
          update_train_iters(args)
          if args.lr_decay_samples is None:
              args.lr_decay_samples = args.train_samples
          lr_decay_steps = args.lr_decay_samples
          wd_incr_steps = args.train_samples
          if args.lr_warmup_fraction is not None:
              lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
          else:
              lr_warmup_steps = args.lr_warmup_samples
      else:
          raise Exception(
              "either train-iters or train-samples should be provided."
          )
  
      opt_param_scheduler = OptimizerParamScheduler(
          optimizer,
          init_lr=args.lr_warmup_init,
          max_lr=args.lr,
          min_lr=args.min_lr,
          lr_warmup_steps=lr_warmup_steps,
          lr_decay_steps=lr_decay_steps,
          lr_decay_style=args.lr_decay_style,
          start_wd=args.start_weight_decay,
          end_wd=args.end_weight_decay,
          wd_incr_steps=wd_incr_steps,
          wd_incr_style=args.weight_decay_incr_style,
          use_checkpoint_opt_param_scheduler=args.use_checkpoint_opt_param_scheduler,
          override_opt_param_scheduler=args.override_opt_param_scheduler,
      )
  
      return opt_param_scheduler
  
imports: import gc
  from datetime import datetime
  import math
  import logging
  import sys
  from .log_handler import CustomHandler
  from .theoretical_memory_usage import report_theoretical_memory
  import time
  import torch
  from megatron import get_args
  from megatron import get_signal_handler
  from megatron import get_timers
  from megatron import get_tensorboard_writer
  from megatron import get_wandb_writer
  from megatron import get_current_global_batch_size
  from megatron import get_num_microbatches
  from megatron import is_last_rank
  from megatron import update_num_microbatches
  from megatron.core import mpu, tensor_parallel
  from megatron.core.utils import get_model_config
  from megatron import print_rank_0
  from megatron import print_rank_last
  from megatron.checkpointing import load_checkpoint
  from megatron.checkpointing import save_checkpoint
  from megatron.model import Float16Module
  from megatron.model import GPTModel
  from megatron.core.distributed import DistributedDataParallel
  from megatron.core.distributed import finalize_model_grads
  from megatron.core.enums import ModelType
  from megatron.optimizer import get_megatron_optimizer
  from megatron.initialize import initialize_megatron
  from megatron.initialize import write_args_to_tensorboard
  from megatron.initialize import set_jit_fusion_options
  from megatron.optimizer_param_scheduler import OptimizerParamScheduler
  from megatron.utils import check_adlr_autoresume_termination
  from megatron.utils import unwrap_model
  from megatron.data.data_samplers import build_pretraining_data_loader
  from megatron.utils import calc_params_l2_norm
  from megatron.core.pipeline_parallel import get_forward_backward_func
  from megatron.utils import report_memory
  from megatron.model.vision.knn_monitor import compute_feature_bank
label_desc: Use image node_id 6 for calling a global function with example usage: get_optimizer_param_scheduler(optimizer) and returns: opt_param_scheduler