From 2e1bb4b65571c1f4c19f2c703e766be27dd290e1 Mon Sep 17 00:00:00 2001 From: Nikolai Kolodziej <7687617+kldzj@users.noreply.github.com> Date: Fri, 21 Nov 2025 02:57:28 +0100 Subject: [PATCH] Use `PYTORCH_ALLOC_CONF` instead of deprecated `PYTORCH_CUDA_ALLOC_CONF` (#32) * Use `PYTORCH_ALLOC_CONF` instead of deprecated `PYTORCH_CUDA_ALLOC_CONF` * style: reformat environment variable check --- src/heretic/main.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index 334f1a1..4446059 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -46,8 +46,11 @@ from .utils import ( def run(): # Enable expandable segments to reduce memory fragmentation on multi-GPU setups. - if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ: - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + if ( + "PYTORCH_ALLOC_CONF" not in os.environ + and "PYTORCH_CUDA_ALLOC_CONF" not in os.environ + ): + os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" # Modified "Pagga" font from https://budavariam.github.io/asciiart-text/ print(f"[cyan]█░█░█▀▀░█▀▄░█▀▀░▀█▀░█░█▀▀[/] v{version('heretic-llm')}")