<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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