code: def _get_best_final_states(self, finished_states: List[StateType]) -> Dict[int, List[StateType]]:
          
          batch_states: Dict[int, List[StateType]] = defaultdict(list)
          for state in finished_states:
              batch_states[state.batch_indices[0]].append(state)
          best_states: Dict[int, List[StateType]] = {}
          for batch_index, states in batch_states.items():
              # The time this sort takes is pretty negligible, no particular need to optimize this
              # yet.  Maybe with a larger beam size...
              finished_to_sort = [(-state.score[0].item(), state) for state in states]
              finished_to_sort.sort(key=lambda x: x[0])
              best_states[batch_index] = [state[1] for state in finished_to_sort[:self._beam_size]]
          return best_states
docs: Returns the best finished states for each batch instance based on model scores. We return
          at most ``self._max_num_decoded_sequences`` number of sequences per instance.