code: def get_batch_size(x: Union[Tensor, Dict, List]) -> int:
      if isinstance(x, Tensor):
          return x.shape[0]
      elif isinstance(x, Dict):
          for key in ("image", "video", "audio"):
              if key in x:
                  return get_batch_size(x[key])
          raise NotImplementedError(f"Invalid dict keys {x.keys()}")
      elif isinstance(x, List):
          return len(x)
      else:
          raise NotImplementedError(f"Invalid type {type(x)}")
  
file_path: ml-cvnets/engine/utils.py
imports: import argparse
  from typing import Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from torch.cuda.amp import autocast
  from utils import logger
  from utils.common_utils import create_directories
  from utils.ddp_utils import is_master
label_desc: Use image node_id 2 for calling a global function with example usage: get_batch_size(x) and returns: x, len, get_batch_size