From 83b45c2a5c448cfc74fd1abf0348ff77f6a05765 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 17:01:28 -0400
Subject: [PATCH] feat: add PYTORCH_CUDA_ALLOC_CONF for memory optimization

---
 train-on-this-server.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/train-on-this-server.sh b/train-on-this-server.sh
index 1d256f4..b744335 100755
--- a/train-on-this-server.sh
+++ b/train-on-this-server.sh
@@ -18,8 +18,9 @@ pip install --upgrade pip
 echo "Installing training dependencies..."
 pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
 
-# Run training
+# Run training with CUDA memory optimization
 echo "Starting training with accelerate..."
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml
 
 echo "Training completed!"
\ No newline at end of file