def _prune_beam(states: List[State],
                    beam_size: int,
                    sort_states: bool = False) -> List[State]:
        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