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
@@ -10,7 +10,7 @@ torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft
 """

 import os
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"

 import wandb
 import torch