Update japanese_bert.py

2023-09-05 20:24:23 +08:00
parent cb80c6550f
commit eee6e42290
1 changed files with 1 additions and 1 deletions
--- a/text/japanese_bert.py
+++ b/text/japanese_bert.py
@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM

 tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")

-def get_bert_feature(text, word2ph):
+def get_bert_feature(text, word2ph, device=None):
    if sys.platform == "darwin" and torch.backends.mps.is_available() and device == "cpu":
        device = "mps"
    if not device: