code: def metainfo_trace(
      gm: torch.fx.GraphModule,
      *args,
      verbose: bool = False,
      unit: str = "MB",
      **kwargs,
  ) -> None:
      """
      MetaInfo tracing API
  
      Given a ``GraphModule`` and a sample input, this API will trace the MetaInfo of a single training cycle,
      and annotate them on ``gm.graph``.
  
      Uses:
          >>> model = ...
          >>> gm = symbolic_trace(model)
          >>> args = ...  # sample input to the ``GraphModule``
          >>> metainfo_trace(gm, *args)
  
      Args:
          gm (torch.fx.GraphModule): The ``GraphModule`` to be annotated with MetaInfo.
          verbose (bool, optional): Whether to show ``MetaInfoProp.summary()`. Defaults to False.
          unit (str, optional): The unit of memory. Defaults to "MB".
  
      Returns:
          torch.fx.GraphModule: The ``GraphModule`` annotated with MetaInfo.
      """
      device = (
          torch.device("cuda:0")
          if torch.cuda.is_available()
          else torch.device("cpu")
      )
      interp = MetaInfoProp(gm.to(device))
      if is_compatible_with_meta():
          from colossalai.fx.profiler import MetaTensor
  
          args = tree_map(
              lambda x: MetaTensor(x, fake_device=device), args
          )
          kwargs = tree_map(
              lambda x: MetaTensor(x, fake_device=device), kwargs
          )
      interp.propagate(*args, **kwargs)
      if verbose:
          interp.summary(unit)
      gm.to("cpu")
      del interp
      return gm
  
imports: from dataclasses import asdict
  from typing import Any, Dict, List, NamedTuple, Tuple
  import torch
  import torch.fx
  from torch.fx.node import Argument, Node, Target
  from torch.utils._pytree import tree_map
  from colossalai.fx._compatibility import compatibility, is_compatible_with_meta
  from colossalai.fx.profiler import GraphInfo, activation_size, calculate_fwd_in, calculate_fwd_out, calculate_fwd_tmp, profile_function, profile_method, profile_module
label_desc: Use image node_id 2 for calling a global function with example usage: metainfo_trace(gm) and returns: gm