code: def batch_tensor_dicts(tensor_dicts: List[Dict[str, torch.Tensor]],
                         remove_trailing_dimension: bool = False) -> Dict[str, torch.Tensor]:
      
      key_to_tensors: Dict[str, List[torch.Tensor]] = defaultdict(list)
      for tensor_dict in tensor_dicts:
          for key, tensor in tensor_dict.items():
              key_to_tensors[key].append(tensor)
      batched_tensors = {}
      for key, tensor_list in key_to_tensors.items():
          batched_tensor = torch.stack(tensor_list)
          if remove_trailing_dimension and all(tensor.size(-1) == 1 for tensor in tensor_list):
              batched_tensor = batched_tensor.squeeze(-1)
          batched_tensors[key] = batched_tensor
      return batched_tensors
docs: Takes a list of tensor dictionaries, where each dictionary is assumed to have matching keys,
      and returns a single dictionary with all tensors with the same key batched together.
  
      Parameters
      ----------
      tensor_dicts : ``List[Dict[str, torch.Tensor]]``
          The list of tensor dictionaries to batch.
      remove_trailing_dimension : ``bool``
          If ``True``, we will check for a trailing dimension of size 1 on the tensors that are being
          batched, and remove it if we find it.