code: def report_memory(name):
      """Simple GPU memory report."""
      mega_bytes = 1024.0 * 1024.0
      string = name + " memory (MB)"
      string += " | allocated: {}".format(
          torch.cuda.memory_allocated() / mega_bytes
      )
      string += " | max allocated: {}".format(
          torch.cuda.max_memory_allocated() / mega_bytes
      )
      string += " | reserved: {}".format(
          torch.cuda.memory_reserved() / mega_bytes
      )
      string += " | max reserved: {}".format(
          torch.cuda.max_memory_reserved() / mega_bytes
      )
      if parallel_state.get_data_parallel_rank() == 0:
          print(
              "[Rank {}] {}".format(
                  torch.distributed.get_rank(), string
              ),
              flush=True,
          )
  
imports: from typing import Optional, List, Union, Tuple
  import torch
  from torch.nn.parallel import DistributedDataParallel
  from apex.multi_tensor_apply import multi_tensor_applier
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.microbatches import build_num_microbatches_calculator
  from apex.transformer.pipeline_parallel._timers import _Timers
label_desc: Use image node_id 23 for calling a global function with example usage: report_memory(name) without return types