def _get_best_final_states(self, finished_states: List[StateType]) -> Dict[int, List[StateType]]:
        """
        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.
        """
        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
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.