code: def tensor_shape_list(obj):
      return pytree_map(
          obj, fn=lambda x: x.shape, process_types=torch.Tensor
      )
  
imports: import argparse
  import os
  import warnings
  from typing import Any, Callable, Tuple, Type, Union
  import torch
  import torch.distributed.rpc
  import torch.multiprocessing
  from torch._C._distributed_rpc import _is_current_rpc_agent_set
  from torch.futures import Future
  from colossalai.initialize import launch
  from colossalai.legacy.pipeline.pipeline_process_group import ppg
label_desc: Use image node_id 3 for calling a global function with example usage: tensor_shape_list(obj) and returns: pytree_map