From 0e71d7ef7c31dfb8e59f7d508d2aa19d7dd1ea81 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Sun, 11 Aug 2024 08:00:56 +0900 Subject: [PATCH] Fix: Specifying the device_map option in extract_bert_feature() --- style_bert_vits2/nlp/chinese/bert_feature.py | 2 +- style_bert_vits2/nlp/english/bert_feature.py | 2 +- style_bert_vits2/nlp/japanese/bert_feature.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/style_bert_vits2/nlp/chinese/bert_feature.py b/style_bert_vits2/nlp/chinese/bert_feature.py index d004593..fe54026 100644 --- 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 diff --git a/style_bert_vits2/nlp/english/bert_feature.py b/style_bert_vits2/nlp/english/bert_feature.py index ede208e..3338c25 100644 --- a/style_bert_vits2/nlp/english/bert_feature.py +++ b/style_bert_vits2/nlp/english/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.EN) + model = bert_models.load_model(Languages.EN, device_map=device) bert_models.transfer_model(Languages.EN, device) style_res_mean = None diff --git a/style_bert_vits2/nlp/japanese/bert_feature.py b/style_bert_vits2/nlp/japanese/bert_feature.py index 0e2d9eb..f426a13 100644 --- a/style_bert_vits2/nlp/japanese/bert_feature.py +++ b/style_bert_vits2/nlp/japanese/bert_feature.py @@ -36,7 +36,7 @@ def extract_bert_feature( if device == "cuda" and not torch.cuda.is_available(): device = "cpu" - model = bert_models.load_model(Languages.JP) + model = bert_models.load_model(Languages.JP, device_map=device) bert_models.transfer_model(Languages.JP, device) style_res_mean = None