Improve memory safety by Release Memory After CV Folds and Optimization Trials

[thomasATbayer]

Release Memory After CV Folds and Optimization Trials

Problem

During cross-validation and hyperparameter optimization, large models (especially GPU-based ones) are kept in memory across folds/trials. This leads to excessive memory consumption and potential OOM errors, particularly when using CUDA-enabled models like CatBoost on GPU or TabPFN.

Solution

Introduced a centralized release_memory() utility in mother/utils.py that:

• Runs Python garbage collection (gc.collect())
• Synchronizes and flushes the CUDA GPU cache when PyTorch is available and a GPU is present
• Is a safe no-op when PyTorch is not installed or no GPU exists

This function is called in two key locations:

1. mother/pipeline_utils.py — mother_cv(): After each CV fold completes, the fold estimator is explicitly deleted and release_memory() is called, freeing both CPU and GPU resources between folds.

2. mother/optimization/core.py — MotherTuner: Replaced the existing bare gc.collect() call with release_memory(), adding GPU cache cleanup to the optimization loop as well.

Changes

File	Change
mother/utils.py	Added release_memory() function with gc.collect() + conditional CUDA cache flush
mother/pipeline_utils.py	Call del val_estimator + release_memory() after each CV fold
mother/optimization/core.py	Replace gc.collect() with release_memory()

Notes

• No new dependencies — PyTorch import is optional and guarded by a try/except at module level
• Fully backward-compatible; no API changes

[thomasATbayer]

def release_memory() -> None:
"""Run garbage collection and, if CUDA is available, flush the GPU cache.

Safe to call unconditionally — it is a no-op when PyTorch is not
installed or no GPU is present.
"""
gc.collect()Expand commentComment on lines R18 to R24Resolved
try:
    import torch
except (ImportError, OSError):
    return
if torch.cuda.is_available():
    torch.cuda.synchronize()
    torch.cuda.empty_cache()