code: def report_memory_usage(message, logger=None, report_cpu=False):
      """Calculate and print RAM usage (in GB)
  
      Args:
          message (str): A prefix message to add in the log.
          logger (:class:`colossalai.logging.DistributedLogger`): The logger used to record memory information.
          report_cpu (bool, optional): Whether to report CPU memory.
  
      Raises:
          EnvironmentError: Raise error if no distributed environment has been initialized.
      """
      if not dist.is_initialized():
          raise EnvironmentError(
              "No distributed environment is initialized"
          )
  
      gpu_allocated = _bytes_to_MB(torch.cuda.memory_allocated())
      gpu_max_allocated = _bytes_to_MB(
          torch.cuda.max_memory_allocated()
      )
      gpu_cached = _bytes_to_MB(torch.cuda.memory_reserved())
      gpu_max_cached = _bytes_to_MB(torch.cuda.max_memory_reserved())
  
      full_log = (
          f"{message}: GPU: allocated {gpu_allocated} MB, max allocated {gpu_max_allocated} MB, "
          + f"cached: {gpu_cached} MB, max cached: {gpu_max_cached} MB"
      )
  
      if report_cpu:
          # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports
          gc.collect()
          vm_stats = psutil.virtual_memory()
          vm_used = _bytes_to_MB(vm_stats.total - vm_stats.available)
          full_log += f", CPU Virtual Memory: used = {vm_used} MB, percent = {vm_stats.percent}%"
  
      if logger is None:
          logger = get_dist_logger()
      logger.info(full_log)
  
      # get the peak memory to report correct data, so reset the counter for the next call
      if hasattr(torch.cuda, "reset_peak_memory_stats"):  # pytorch 1.4+
          torch.cuda.reset_peak_memory_stats()
  
file_path: colossalai/colossalai/legacy/utils/memory.py
imports: import gc
  from collections import namedtuple
  import psutil
  import torch
  import torch.distributed
  from packaging import version
  from colossalai.legacy.core import global_context
  from colossalai.logging import get_dist_logger
  from colossalai.utils import get_current_device
label_desc: Use image node_id 3 for calling a global function with example usage: report_memory_usage(message, logger, report_cpu) without return types