code: def combine_initial_dims(tensor: torch.Tensor) -> torch.Tensor:
      
      if tensor.dim() <= 2:
          return tensor
      else:
          return tensor.view(-1, tensor.size(-1))
docs: Given a (possibly higher order) tensor of ids with shape
      (d1, ..., dn, sequence_length)
      Return a view that's (d1 * ... * dn, sequence_length).
      If original tensor is 1-d or 2-d, return it as is.