Restore gc empty_cache in train_ms_jp

This commit is contained in:
litagin02
2024-03-13 10:25:22 +09:00
parent 30d9f2e258
commit 28fd715cb8

View File

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