code: 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
docs: 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.