From 28fd715cb89b11711b128521f3b046c89e6fcca1 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Wed, 13 Mar 2024 10:25:22 +0900
Subject: [PATCH] Restore gc empty_cache in train_ms_jp

---
 train_ms_jp_extra.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py
index 722a521..7b2cb42 100644
--- a/train_ms_jp_extra.py
+++ b/train_ms_jp_extra.py
@@ -1,5 +1,6 @@
 import argparse
 import datetime
+import gc
 import os
 import platform
 
@@ -1008,9 +1009,9 @@ def train_and_evaluate(
                 )
             )
             pbar.update()
-    # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
-    # gc.collect()
-    # torch.cuda.empty_cache()
+
+    gc.collect()
+    torch.cuda.empty_cache()
     if pbar is None and rank == 0:
         logger.info(f"====> Epoch: {epoch}, step: {global_step}")