Refactor: Use bert_models.transfer_model()

2024-08-10 16:18:14 +09:00
parent f84d2ed48b
commit bdf20e8911
5 changed files with 13 additions and 4 deletions
--- a/style_bert_vits2/nlp/japanese/bert_feature.py
+++ b/style_bert_vits2/nlp/japanese/bert_feature.py
@@ -36,7 +36,8 @@ def extract_bert_feature(

    if device == "cuda" and not torch.cuda.is_available():
        device = "cpu"
-    model = bert_models.load_model(Languages.JP).to(device)  # type: ignore
+    model = bert_models.load_model(Languages.JP)
+    bert_models.transfer_model(Languages.JP, device)

    style_res_mean = None
    with torch.no_grad():