From 649b4fb9c9f40fb7c4e307248de2f9f34a9537b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=AC=B8=E5=98=BF=EF=BC=81?=
 <50474364+GhostZey@users.noreply.github.com>
Date: Sun, 10 Dec 2023 21:56:13 +0800
Subject: [PATCH] Update train_ms.py (#225)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

在每次epoch结束后回收显存以避免显存溢出的问题
---
 train_ms.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/train_ms.py b/train_ms.py
index 659669f..30fd315 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -13,6 +13,7 @@ import logging
 from config import config
 import argparse
 import datetime
+import gc
 
 logging.getLogger("numba").setLevel(logging.WARNING)
 import commons
@@ -590,7 +591,8 @@ def train_and_evaluate(
                     )
 
         global_step += 1
-
+    gc.collect()
+    torch.cuda.empty_cache()
     if rank == 0:
         logger.info("====> Epoch: {}".format(epoch))