Fix: Add missing torch.cuda.empty_cache() (fishaudio#161)

leitao97 · Nov 8, 2023 · 6e66fec · 6e66fec
1 parent 6a77c8b
commit 6e66fec
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/infer.py b/infer.py
@@ -202,4 +202,6 @@ def infer(
             .numpy()
         )
         del x_tst, tones, lang_ids, bert, x_tst_lengths, speakers, ja_bert, en_bert
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         return audio
diff --git a/oldVersion/V101/__init__.py b/oldVersion/V101/__init__.py
@@ -70,4 +70,6 @@ def infer(
             .numpy()
         )
         del x_tst, tones, lang_ids, bert, x_tst_lengths, speakers
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         return audio