code: def _get_multiplier(stage: WarmupStage, iter: int) -> float:
      multiplier = 1.0
      if stage.policy == WarmupPolicy.LINEAR:
          multiplier = (
              stage.value + (1.0 - stage.value) * iter / stage.max_iters
          )
      elif stage.policy == WarmupPolicy.CONSTANT:
          multiplier = stage.value
      elif stage.policy == WarmupPolicy.POLY:
          multiplier = math.pow(
              1 - iter / stage.decay_iters, stage.value
          )
      elif stage.policy == WarmupPolicy.STEP:
          multiplier = math.pow(stage.value, iter // stage.decay_iters)
      elif stage.policy == WarmupPolicy.INVSQRT:
          multiplier = 1.0 / math.sqrt(iter)
      return multiplier * stage.lr_scale
  
file_path: torchrec/torchrec/optim/warmup.py
imports: import logging
  import math
  from dataclasses import dataclass
  from enum import Enum, unique
  from typing import Any, List, Tuple
  import torch
  from torchrec.optim.keyed import KeyedOptimizer, OptimizerWrapper
label_desc: Use image node_id 2 for calling a global function with example usage: _get_multiplier(stage, iter) and returns: unknown