code: def colo_model_data_tensor_move(
      src_t: Union[StatefulTensor, torch.Tensor],
      tgt_t: Union[StatefulTensor, torch.Tensor],
  ) -> None:
      """
      A colossal API for model data tensor move.
      The src and target tensors could be resident on both CPU and GPU.
  
      NOTE() The source tensor payload will be removed after this function.
  
      The function will record the communication volume between CPU and GPU.
      Args:
          src_t (Union[StatefulTensor, torch.Tensor]): source tensor
          tgt_t (Union[StatefulTensor, torch.Tensor]): target tensor
      """
      if isinstance(src_t, StatefulTensor):
          src_t_payload = src_t.payload
      else:
          src_t_payload = src_t.data
      src_dev = src_t_payload.device
  
      if isinstance(tgt_t, StatefulTensor):
          tgt_t_payload = tgt_t.payload
      else:
          tgt_t_payload = tgt_t.data
  
      tgt_t_payload.copy_(src_t_payload)
  
      # remove payload of src_t
      if isinstance(src_t, StatefulTensor):
          src_t.set_null()
      else:
          src_t.data = torch.empty(
              0, device=src_dev, dtype=src_t_payload.dtype
          )
  
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 5 for calling a global function with example usage: colo_model_data_tensor_move(src_t, tgt_t) without return types