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