Update train_ms.py (#225)

在每次epoch结束后回收显存以避免显存溢出的问题

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))