Fix: Ensure memory is cleared on unload

2024-11-10 07:06:20 +09:00
parent 9723f75f99
commit 20b0b0908b
2 changed files with 21 additions and 7 deletions
--- a/style_bert_vits2/nlp/bert_models.py
+++ b/style_bert_vits2/nlp/bert_models.py
@@ -237,9 +237,9 @@ def unload_model(language: Languages) -> None:

    if language in __loaded_models:
        del __loaded_models[language]
-        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
+        gc.collect()
        logger.info(f"Unloaded the {language.name} BERT model")


@@ -256,8 +256,6 @@ def unload_tokenizer(language: Languages) -> None:
    if language in __loaded_tokenizers:
        del __loaded_tokenizers[language]
        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
        logger.info(f"Unloaded the {language.name} BERT tokenizer")