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