code: def get_batch_lengths(batch):
      lengths = []
      pytree_map(
          batch,
          fn=lambda x: lengths.append(len(x)),
          process_types=torch.Tensor,
      )
      return lengths
  
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 4 for calling a global function with example usage: get_batch_lengths(batch) and returns: lengths