class_docstr: Stateful Tensor Manager, inspired from PatrickStar
  
  PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
  https://arxiv.org/abs/2108.05818
class_docstr_tok: ["Stateful","Tensor","Manager",",","inspired","from","PatrickStar","PatrickStar",":","Parallel","Training","of","Pre-trained","Models","via","Chunk-based","Memory","Management","https",":","\/\/arxiv.org\/abs\/2108.05818"]
code: def _get_layout_info(self, compute_idx: int, warmup: bool):
      move_to_cuda_tensor_list = []
      hold_cuda_tensor_list = []
      for tensor in self._stateful_tensor_list:
          if tensor.state == TensorState.FREE:
              continue
  
          if tensor.device.type == "cuda":
              if tensor.state in [
                  TensorState.HOLD,
                  TensorState.HOLD_AFTER_BWD,
                  TensorState.HOLD_AFTER_FWD,
              ]:
                  hold_cuda_tensor_list.append(tensor)
          elif tensor.device.type == "cpu":
              if tensor.state == TensorState.COMPUTE:
                  move_to_cuda_tensor_list.append(tensor)
          else:
              raise RuntimeError
      return move_to_cuda_tensor_list, hold_cuda_tensor_list
  
file_path: colossalai/colossalai/legacy/zero/gemini/stateful_tensor_mgr.py
imports: import functools
  import types
  from time import time
  from typing import List
  from colossalai.utils.device import get_current_device
  from .stateful_tensor import StatefulTensor, TensorState
  from .tensor_placement_policy import TensorPlacementPolicy
  from .tensor_utils import colo_model_data_tensor_move_inline
label_desc: Use image node_id 8 for calling the StatefulTensorMgr obj's underlying member method code with example usage: obj._get_layout_info(compute_idx, warmup) and returns: move_to_cuda_tensor_list,  hold_cuda_tensor_list