添加device传参

text/japanese_bert.py

@@ -1,20 +1,13 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-device = torch.device(
-        "cuda"
-        if torch.cuda.is_available()
-        else (
-            "mps"
-            if sys.platform == "darwin" and torch.backends.mps.is_available()
-            else "cpu"
-        )
-    )
 
 tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")
-model = AutoModelForMaskedLM.from_pretrained("./bert/bert-base-japanese-v3").to(device)
 
 def get_bert_feature(text, word2ph):
+    if sys.platform == "darwin" and torch.backends.mps.is_available() and device == "cpu"
+        device == "mps"
+    model = AutoModelForMaskedLM.from_pretrained("./bert/bert-base-japanese-v3").to(device)
     with torch.no_grad():
         inputs = tokenizer(text, return_tensors='pt')
         for i in inputs: