code: def get_batch_size(batch: Any) -> int:
      """Get the batch size (size of dimension-0) of the first tensor in the batch.
  
      Args:
          batch (Any): Batch to be inspected.
  
      Raises:
          RuntimeError: If no tensor is found in the batch.
  
      Returns:
          int: Batch size.
      """
      data_list, _ = tree_flatten(batch)
      for data in data_list:
          if isinstance(data, torch.Tensor):
              return data.size(0)
      raise RuntimeError("No tensor found in the batch")
  
file_path: colossalai/colossalai/pipeline/schedule/_utils.py
imports: from collections import OrderedDict
  from typing import Any, List, Optional, Tuple
  import torch
  import torch.cuda
  from torch.nn import Module
  from torch.utils._pytree import SUPPORTED_NODES, TreeSpec, _register_pytree_node, tree_flatten, tree_map, tree_unflatten
label_desc: Use image node_id 6 for calling a global function with example usage: get_batch_size(batch) and returns: data