code: def memory_usage(self) -> Dict[str, int]:
      cuda_memory = 0
      cpu_memory = 0
  
      if self.chunk_temp is not None:
          # this chunk is not closed
          if (
              self.chunk_temp.device.type == "cuda"
              or self.chunk_temp.device.type == "npu"
          ):
              cuda_memory += self.chunk_mem
          else:
              cpu_memory += self.chunk_mem
      else:
          if self.is_gathered:
              cuda_memory += self.chunk_mem
          if self.cuda_shard is not None:
              cuda_memory += self.shard_mem
          if self.cpu_shard is not None:
              cpu_memory += self.shard_mem
  
      return dict(cuda=cuda_memory, cpu=cpu_memory)
  
imports: from dataclasses import dataclass
  from enum import Enum
  from typing import Dict, List, Optional
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
label_desc: Use image node_id 2 for calling the Chunk obj's underlying member method code with example usage: obj.memory_usage() and returns: dict