code: def wrap_module(
      BaseT: Type[nn.Module], DestT: Type[nn.Module]
  ) -> Callable[[nn.Module], Optional[nn.Module]]:
      """
      Generic function generator to replace BaseT module with DestT wrapper.
      Args:
          BaseT : module type to replace
          DestT : destination module type
      Returns:
          swap function to replace BaseT module with DestT
      """
  
      def expansion_fn(mod: nn.Module) -> Optional[nn.Module]:
          out = DestT(mod)
          return out
  
      return expansion_fn
  
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 13 for calling a global function with example usage: wrap_module(BaseT, DestT) and returns: expansion_fn, out