code: def colo_device_memory_capacity(device: torch.device) -> int:
      """
      Get the capacity of the memory of the device
  
      Args:
          device (torch.device): a device
  
      Returns:
          int: size in byte
      """
      # TODO: add NPU support
      assert isinstance(device, torch.device)
      if device.type == "cpu":
          # In the context of 1-CPU-N-GPU, the memory capacity of the current process is 1/N overall CPU memory.
          return colo_get_cpu_memory_capacity() // dist.get_world_size()
      if device.type == "cuda":
          return (
              torch.cuda.get_device_properties(
                  get_current_device()
              ).total_memory
              * _GLOBAL_CUDA_MEM_FRACTION
          )
  
file_path: colossalai/colossalai/utils/memory.py
imports: from collections import namedtuple
  import psutil
  import torch
  import torch.distributed
  from colossalai.utils import get_current_device
label_desc: Use image node_id 2 for calling a global function with example usage: colo_device_memory_capacity(device) and returns: unknown, unknown