Fix: Specifying the device_map option in extract_bert_feature()

2024-08-11 08:00:56 +09:00
parent 80dd6dbc22
commit 0e71d7ef7c
3 changed files with 3 additions and 3 deletions
--- a/style_bert_vits2/nlp/chinese/bert_feature.py
+++ b/style_bert_vits2/nlp/chinese/bert_feature.py
@@ -29,7 +29,7 @@ def extract_bert_feature(

    if device == "cuda" and not torch.cuda.is_available():
        device = "cpu"
-    model = bert_models.load_model(Languages.ZH)
+    model = bert_models.load_model(Languages.ZH, device_map=device)
    bert_models.transfer_model(Languages.ZH, device)

    style_res_mean = None