code: def make_lr_schedule(
      args: argparse.Namespace, optimizer: torch.optim.Optimizer
  ) -> np.ndarray:
      """Helper function to return a learning rate scheduler for CRE-stereo"""
      if args.decay_after_steps < args.warmup_steps:
          raise ValueError(
              f"decay_after_steps: {args.function} must be greater than warmup_steps: {args.warmup_steps}"
          )
  
      warmup_steps = args.warmup_steps if args.warmup_steps else 0
      flat_lr_steps = (
          args.decay_after_steps - warmup_steps
          if args.decay_after_steps
          else 0
      )
      decay_lr_steps = args.total_iterations - flat_lr_steps
  
      max_lr = args.lr
      min_lr = args.min_lr
  
      schedulers = []
      milestones = []
  
      if warmup_steps > 0:
          if args.lr_warmup_method == "linear":
              warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                  optimizer,
                  start_factor=args.lr_warmup_factor,
                  total_iters=warmup_steps,
              )
          elif args.lr_warmup_method == "constant":
              warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
                  optimizer,
                  factor=args.lr_warmup_factor,
                  total_iters=warmup_steps,
              )
          else:
              raise ValueError(
                  f"Unknown lr warmup method {args.lr_warmup_method}"
              )
          schedulers.append(warmup_lr_scheduler)
          milestones.append(warmup_steps)
  
      if flat_lr_steps > 0:
          flat_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
              optimizer, factor=max_lr, total_iters=flat_lr_steps
          )
          schedulers.append(flat_lr_scheduler)
          milestones.append(flat_lr_steps + warmup_steps)
  
      if decay_lr_steps > 0:
          if args.lr_decay_method == "cosine":
              decay_lr_scheduler = (
                  torch.optim.lr_scheduler.CosineAnnealingLR(
                      optimizer, T_max=decay_lr_steps, eta_min=min_lr
                  )
              )
          elif args.lr_decay_method == "linear":
              decay_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                  optimizer,
                  start_factor=max_lr,
                  end_factor=min_lr,
                  total_iters=decay_lr_steps,
              )
          elif args.lr_decay_method == "exponential":
              decay_lr_scheduler = (
                  torch.optim.lr_scheduler.ExponentialLR(
                      optimizer,
                      gamma=args.lr_decay_gamma,
                      last_epoch=-1,
                  )
              )
          else:
              raise ValueError(
                  f"Unknown lr decay method {args.lr_decay_method}"
              )
          schedulers.append(decay_lr_scheduler)
  
      scheduler = torch.optim.lr_scheduler.SequentialLR(
          optimizer, schedulers, milestones=milestones
      )
      return scheduler
  
imports: import argparse
  import os
  import warnings
  from pathlib import Path
  from typing import List, Union
  import numpy
  import torch
  import torch.distributed
  import torchvision.models.optical_flow
  import torchvision.prototype.models.depth.stereo
  import utils
  import visualization
  from parsing import make_dataset, make_eval_transform, make_train_transform, VALID_DATASETS
  from torch import nn
  from torchvision.transforms.functional import get_dimensions, InterpolationMode, resize
  from utils.metrics import AVAILABLE_METRICS
  from utils.norm import freeze_batch_norm
label_desc: Use image node_id 2 for calling a global function with example usage: make_lr_schedule(args, optimizer) and returns: scheduler