def cleanup_memory():
    """Cleanup memory by running garbage collection and emptying CUDA cache."""
    import torch

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
Cleanup memory by running garbage collection and emptying CUDA cache.
olive/common/utils.py