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
This commit is contained in:
Nikolai Kolodziej
2025-11-21 02:57:28 +01:00
committed by GitHub
parent af02bc6ece
commit 2e1bb4b655
+5 -2
View File
@@ -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')}")