code: def pack_return_tensors(return_tensors):
      output, label = tuple(zip(*return_tensors))
      if isinstance(output[0], torch.Tensor):
          output = torch.cat(output, dim=0)
      elif isinstance(output[0], (list, tuple)):
          output = tuple(
              torch.cat(tensors, dim=0) for tensors in zip(*output)
          )
      else:
          raise TypeError(
              f"Output of model must be tensor or list/tuple of tensors"
          )
      if isinstance(label[0], torch.Tensor):
          label = torch.cat(label, dim=0)
      else:
          merged_label = {k: [] for k in label[0].keys()}
          for d in label:
              for k, v in d.items():
                  merged_label[k].append(v)
          label = {
              k: torch.cat(v, dim=0) for k, v in merged_label.items()
          }
      return output, label
  
imports: import inspect
  from typing import Callable, List, Tuple, Union
  import torch.cuda
  import colossalai.legacy.communication
  from colossalai.legacy.amp.naive_amp import NaiveAMPModel
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.utils import switch_virtual_pipeline_parallel_rank
  from colossalai.logging import get_dist_logger
  from colossalai.utils.device import get_current_device
  from ._base_schedule import BaseSchedule
label_desc: Use image node_id 2 for calling a global function with example usage: pack_return_tensors(return_tensors) and returns: output,  label