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