def METHOD_NAME(device: torch.device) -> int:
    """
    Get the capacity of the memory of the device
    Args:
        device (torch.device): a device
    Returns:
        int: size in byte
    """
    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() / gpc.num_processes_on_current_node
    if device.type == 'cuda':
        return torch.cuda.get_device_properties(get_current_device()).total_memory * _GLOBAL_CUDA_MEM_FRACTION