def get_batch_token_embeddings(layer_hidden_states, attention_mask, rm_special_tokens=False):
    """ remove padding and special tokens
    Args:
        layer_hidden_states: (N, L, D)
        attention_mask: (N, L) with 1 indicate valid bits, 0 pad bits
        rm_special_tokens: bool, whether to remove special tokens, this is different for different model_type
            1) a RoBERTa sequence has the following format: <s> X </s>
    return:
        list(np.ndarray), each ndarray is (L_sentence, D), where L_sentence <= L
    """
    valid_lengths = attention_mask.sum(1).long().tolist()  # (N, )
    layer_hidden_states = layer_hidden_states.cpu().numpy()
    embeddings = [e[1:vl-1] if rm_special_tokens else e[:vl]
                  for e, vl in zip(layer_hidden_states, valid_lengths)]
    return embeddings