Skip to content

Commit cb36de1

Browse files
fix gpu memory growth (#14044)
1 parent 5cf3ac5 commit cb36de1

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

tools/program.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from __future__ import print_function
1818

1919
import os
20+
import gc
2021
import sys
2122
import platform
2223
import yaml
@@ -495,6 +496,7 @@ def train(
495496
model,
496497
os.path.join(save_model_dir, prefix, "inference"),
497498
)
499+
gc.collect()
498500
model_info = {"epoch": epoch, "metric": best_model_dict}
499501
else:
500502
model_info = None
@@ -542,6 +544,7 @@ def train(
542544
prefix = "latest"
543545
if uniform_output_enabled:
544546
export(config, model, os.path.join(save_model_dir, prefix, "inference"))
547+
gc.collect()
545548
model_info = {"epoch": epoch, "metric": best_model_dict}
546549
else:
547550
model_info = None
@@ -570,6 +573,7 @@ def train(
570573
prefix = "iter_epoch_{}".format(epoch)
571574
if uniform_output_enabled:
572575
export(config, model, os.path.join(save_model_dir, prefix, "inference"))
576+
gc.collect()
573577
model_info = {"epoch": epoch, "metric": best_model_dict}
574578
else:
575579
model_info = None

0 commit comments

Comments
 (0)