pytorch - BERT トークン化に CPU の代わりに GPU を使用する方法はありますか?

Question

私は文の大規模なデータセット (230 万行、65 億 3000 万語) に対して BERT トークナイザーを使用しています。

#creating a BERT tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', 
                                          do_lower_case=True)

#encoding the data using our tokenizer
encoded_dict = tokenizer.batch_encode_plus(
    df[df.data_type=='train'].comment.values, 
    add_special_tokens=True, 
    return_attention_mask=True, 
    pad_to_max_length=True, 
    max_length=256, 
    return_tensors='pt'
)

そのままCPUで動作し、1コアのみで動作します。並列化を試みましたが、16 コアの CPU では処理速度が 16 倍しか向上せず、完全なデータセットをトークン化する場合でも長時間実行できます。

GPU で実行する方法や、他の方法で高速化する方法はありますか?

編集:高速トークナイザーも使用してみました:

#creating a BERT tokenizer
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased', 
                                          do_lower_case=True)

次に、出力を batch_encode_plus に渡します。

#encoding the data using our tokenizer
encoded_dict = tokenizer.batch_encode_plus(
    df[df.data_type=='train'].comment.values, 
    add_special_tokens=True, 
    return_attention_mask=True, 
    pad_to_max_length=True, 
    max_length=256, 
    return_tensors='pt'
)

しかし、batch_encode_plus は次のエラーを返します。

TypeError: batch_text_or_text_pairs はリストでなければなりません (取得 <class 'numpy.ndarray'>)