code: def replace_modules(
      model: nn.Module,
      expansions: Dict[
          str, Callable[[nn.Module], Optional[nn.Module]]
      ] = None,
  ) -> nn.Module:
      """
      Top-level function to replace modules in model, specified by class name with a desired replacement.
      NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
      Args:
          model : top level module
          expansions : replacement dictionary: module class name -> replacement function generator
      Returns:
          model, possibly modified in-place
      """
      mapping: Dict[str, nn.Module] = {}
      for name, m in model.named_modules():
          m_type = type(m).__name__
          if m_type in expansions:
              swapped = expansions[m_type](m)
              if swapped:
                  mapping[name] = swapped
      if len(mapping) > 0:
          logging.info(f"Swapped {len(mapping)} modules")
      swap_modules(model, mapping)
      return model
  
imports: import os
  from contextlib import nullcontext
  from enum import Enum
  from typing import Callable, Dict, Optional, Type
  import onnx
  import torch
  import torch.nn
  import torch.nn.functional
  from nemo.utils import CastToFloat, CastToFloatAll, logging
label_desc: Use image node_id 15 for calling a global function with example usage: replace_modules(model, expansions) and returns: model