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
@@ -11,7 +11,7 @@ torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device_batch_

 from collections import deque
 import os
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
 import time
 import wandb
 import torch