update to new os env var to get rid of deprecation warning

2025-12-28 03:32:46 +00:00
parent e1770a3061
commit 2874eda59a
3 changed files with 3 additions and 3 deletions
@@ -12,7 +12,7 @@ python -m scripts.base_train --depth=4 --max_seq_len=512 --device_batch_size=1 -
 """

 import os
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
 import time
 from contextlib import nullcontext