code: def memory_status(msg, print_rank=-1, reset_max=False):
      global mem_alloced, mem_cached
  
      rank = dist.get_rank()
      if print_rank != -1 and rank != print_rank:
          return
  
      torch.cuda.synchronize()
  
      if reset_max:
          torch.cuda.reset_max_memory_cached()
          torch.cuda.reset_max_memory_allocated()
  
      new_alloced = torch.cuda.memory_allocated()
      new_cached = torch.cuda.memory_cached()
  
      delta_alloced = new_alloced - mem_alloced
      delta_cached = new_cached - mem_cached
  
      mem_cached = new_cached
      mem_alloced = new_alloced
  
      max_alloced = torch.cuda.max_memory_allocated()
      max_cached = torch.cuda.max_memory_cached()
  
      # convert to GB for printing
      new_alloced /= 1024**3
      new_cached /= 1024**3
      delta_alloced /= 1024**3
      delta_cached /= 1024**3
      max_alloced /= 1024**3
      max_cached /= 1024**3
  
      print(
          f"RANK={rank} MEMSTATS",
          msg,
          f"device={torch.cuda.current_device()} "
          f"current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) "
          f"current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)",
      )
  
imports: from collections.abc import Iterable
  from deepspeed.moe.utils import is_moe_param
  import os
  import psutil
  import gc
  from math import sqrt
  from math import floor
  from bisect import bisect_left
  import torch
  from torch import inf
  from deepspeed import comm
  from deepspeed.utils import groups, logger
  from deepspeed.runtime.constants import PIPE_REPLICATED
  from numpy import prod
label_desc: Use image node_id 19 for calling a global function with example usage: memory_status(msg, print_rank, reset_max) without return types