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