fix gpu memory growth (#14044)

PaddlePaddle · Oct 18, 2024 · cb36de1 · cb36de1
1 parent 5cf3ac5
commit cb36de1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tools/program.py b/tools/program.py
@@ -17,6 +17,7 @@
 from __future__ import print_function
 
 import os
+import gc
 import sys
 import platform
 import yaml
@@ -495,6 +496,7 @@ def train(
                             model,
                             os.path.join(save_model_dir, prefix, "inference"),
                         )
+                        gc.collect()
                         model_info = {"epoch": epoch, "metric": best_model_dict}
                     else:
                         model_info = None
@@ -542,6 +544,7 @@ def train(
             prefix = "latest"
             if uniform_output_enabled:
                 export(config, model, os.path.join(save_model_dir, prefix, "inference"))
+                gc.collect()
                 model_info = {"epoch": epoch, "metric": best_model_dict}
             else:
                 model_info = None
@@ -570,6 +573,7 @@ def train(
             prefix = "iter_epoch_{}".format(epoch)
             if uniform_output_enabled:
                 export(config, model, os.path.join(save_model_dir, prefix, "inference"))
+                gc.collect()
                 model_info = {"epoch": epoch, "metric": best_model_dict}
             else:
                 model_info = None