def METHOD_NAME(self, state, i, new_id=None):
    """Select state with relative ids in the main beam search.
    Args:
        state: Decoder state for prefix tokens
        i (int): Index to select a state in the main beam search
        new_id (int): New label id to select a state if necessary
    Returns:
        state: pruned state
    """
    if type(state) == tuple:
        if len(state) == 2:  # for CTCPrefixScore
            sc, st = state
            return sc[i], st[i]
        else:  # for CTCPrefixScoreTH (need new_id > 0)
            r, log_psi, f_min, f_max, scoring_idmap = state
            s = log_psi[i, new_id].expand(log_psi.size(1))
            if scoring_idmap is not None:
                return r[:, :, i, scoring_idmap[i, new_id]], s, f_min, f_max
            else:
                return r[:, :, i, new_id], s, f_min, f_max
    return None if state is None else state[i]