Amended quickstart.py for better memory allocation

app.py

@@ -2,7 +2,6 @@ import os
 import gradio as gr
 
 api_key = os.getenv("HF_TOKEN")
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 def start_training():
     os.system("python quickstart.py --model_name_or_path Qwen/Qwen3-0.6B --train_data_dir trl-lib/Capybara --api_key HF_TOKEN" )
     return "Evaluation completed!"

quickstart.py

@@ -1,7 +1,9 @@
+import os
 from trl import SFTTrainer
 from datasets import load_dataset
 
 dataset = load_dataset("trl-lib/Capybara", split="train")
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
 trainer = SFTTrainer(
     model="Qwen/Qwen3-0.6B",