code: def colo_tensor_mem_usage(
      tensor: Union[torch.Tensor, StatefulTensor]
  ) -> Tuple[int, int]:
      if isinstance(tensor, StatefulTensor):
          t = tensor.payload
      elif isinstance(tensor, torch.Tensor):
          t = tensor
      else:
          return 0, 0
  
      cuda_use, cpu_use = 0, 0
  
      mem_use = t.storage().size() * t.element_size()
      if t.device.type == "cuda":
          cuda_use += mem_use
      elif t.device.type == "cpu":
          cpu_use += mem_use
  
      return cuda_use, cpu_use
  
file_path: ColossalAI/colossalai/legacy/zero/gemini/tensor_utils.py
imports: from typing import Tuple, Union
  import torch
  from .stateful_tensor import StatefulTensor
label_desc: Use image node_id 4 for calling a global function with example usage: colo_tensor_mem_usage(tensor) and returns: cuda_use,  cpu_use, int,  int