code: def colo_model_mem_usage(model: torch.nn.Module) -> Tuple[int, int]:
      """
      Trace the model memory usage.
      Args:
          model (torch.nn.Module): a torch model
  
      Returns:
          Tuple[int, int]: cuda memory usage in Byte, cpu memory usage in Byte
      """
      if model is None:
          return 0, 0
  
      def _get_tensor_mem_use(t: Optional[torch.Tensor]):
          if t is None:
              return 0, 0
          assert isinstance(t, torch.Tensor)
          _cpu_mem_usage, _cuda_mem_usage = 0, 0
          if t.device.type == "cpu":
              _cpu_mem_usage += t.numel() * t.element_size()
          elif t.device.type == "cuda":
              _cuda_mem_usage += t.numel() * t.element_size()
          return _cuda_mem_usage, _cpu_mem_usage
  
      cuda_mem_usage = 0
      cpu_mem_usage = 0
      for param in model.parameters():
          if hasattr(param, "colo_attr"):
              t_cuda, t_cpu = param.colo_attr.get_memory_usage()
              cuda_mem_usage += t_cuda
              cpu_mem_usage += t_cpu
          else:
              t_cuda, t_cpu = _get_tensor_mem_use(param.data)
              cuda_mem_usage += t_cuda
              cpu_mem_usage += t_cpu
              t_cuda, t_cpu = _get_tensor_mem_use(param.grad)
              cuda_mem_usage += t_cuda
              cpu_mem_usage += t_cpu
  
      return cuda_mem_usage, cpu_mem_usage
  
label_desc: Use image node_id 2 for calling a global function with example usage: colo_model_mem_usage(model) and returns: cuda_mem_usage,  cpu_mem_usage, int,  int, _cuda_mem_usage,  _cpu_mem_usage, int,  int