This method can be used to prune the set of unfinished states on a beam or finished states
        at the end of search. In the former case, the states need not be sorted because the all come
        from the same decoding step, which does the sorting. However, if the states are finished and
        this method is called at the end of the search, they need to be sorted because they come
        from different decoding steps.

def _prune_beam(states: List[State],
                    beam_size: int,
                    sort_states: bool = False) -> List[State]:
        """
        This method can be used to prune the set of unfinished states on a beam or finished states
        at the end of search. In the former case, the states need not be sorted because the all come
        from the same decoding step, which does the sorting. However, if the states are finished and
        this method is called at the end of the search, they need to be sorted because they come
        from different decoding steps.
        """
        states_by_batch_index: Dict[int, List[State]] = defaultdict(list)
        for state in states:
            assert len(state.batch_indices) == 1
            batch_index = state.batch_indices[0]
            states_by_batch_index[batch_index].append(state)
        pruned_states = []
        for _, instance_states in states_by_batch_index.items():
            if sort_states:
                scores = torch.cat([state.score[0].view(-1) for state in instance_states])
                _, sorted_indices = scores.sort(-1, descending=True)
                sorted_states = [instance_states[i] for i in sorted_indices.detach().cpu().numpy()]
                instance_states = sorted_states
            for state in instance_states[:beam_size]:
                pruned_states.append(state)
        return pruned_states