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():
|
def run():
|
||||||
# Enable expandable segments to reduce memory fragmentation on multi-GPU setups.
|
# Enable expandable segments to reduce memory fragmentation on multi-GPU setups.
|
||||||
if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
|
if (
|
||||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
"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/
|
# Modified "Pagga" font from https://budavariam.github.io/asciiart-text/
|
||||||
print(f"[cyan]█░█░█▀▀░█▀▄░█▀▀░▀█▀░█░█▀▀[/] v{version('heretic-llm')}")
|
print(f"[cyan]█░█░█▀▀░█▀▄░█▀▀░▀█▀░█░█▀▀[/] v{version('heretic-llm')}")
|
||||||
|
|||||||
Reference in New Issue
Block a user