code: def payload_mem(self) -> int:
      # sanity check
      assert self.chunk_temp is None
  
      if self.is_gathered:
          return self.chunk_mem
      else:
          return self.shard_mem
  
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 5 for calling the Chunk obj's underlying member method code with example usage: obj.payload_mem() and returns: self, self