code: def get_batch_size(batch: Union[Dict, torch.Tensor]) -> int:
      
      if isinstance(batch, torch.Tensor):
          return batch.size(0) # type: ignore
      elif isinstance(batch, Dict):
          return get_batch_size(next(iter(batch.values())))
      else:
          return 0
docs: Returns the size of the batch dimension. Assumes a well-formed batch,
      returns 0 otherwise.