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:
committed by
GitHub
parent
af02bc6ece
commit
2e1bb4b655
+5
-2
@@ -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')}")
|
||||
|
||||
Reference in New Issue
Block a user