code: def flattened_index_select(target: torch.Tensor,
                             indices: torch.LongTensor) -> torch.Tensor:
      
      if indices.dim() != 2:
          raise ConfigurationError("Indices passed to flattened_index_select had shape {} but "
                                   "only 2 dimensional inputs are supported.".format(indices.size()))
      # Shape: (batch_size, set_size * subset_size, embedding_size)
      flattened_selected = target.index_select(1, indices.view(-1))
  
      # Shape: (batch_size, set_size, subset_size, embedding_size)
      selected = flattened_selected.view(target.size(0), indices.size(0), indices.size(1), -1)
      return selected
docs: The given ``indices`` of size ``(set_size, subset_size)`` specifies subsets of the ``target``
      that each of the set_size rows should select. The `target` has size
      ``(batch_size, sequence_length, embedding_size)``, and the resulting selected tensor has size
      ``(batch_size, set_size, subset_size, embedding_size)``.
  
      Parameters
      ----------
      target : ``torch.Tensor``, required.
          A Tensor of shape (batch_size, sequence_length, embedding_size).
      indices : ``torch.LongTensor``, required.
          A LongTensor of shape (set_size, subset_size). All indices must be < sequence_length
          as this tensor is an index into the sequence_length dimension of the target.
  
      Returns
      -------
      selected : ``torch.Tensor``, required.
          A Tensor of shape (batch_size, set_size, subset_size, embedding_size).