added forced release of GPU memory after job is done
import copy
+import gc
import os
import torch.multiprocessing
import concurrent.futures
@@ -213,6 +214,10 @@ def _run_train_job(sicnk, device=None):
**best,
)
+ # force releasing the GPU memory of the job to avoid memory leakage
+ del job
+ gc.collect()
+
return (train_job_index, best, best_metric)
except BaseException as e:
search_job.config.log(
