Update train.py

train.py

@@ -6,10 +6,9 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer, Train
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["HF_DATASETS_CACHE"] = "/app/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
-save_dir = "/home/user/t5-finetuned"  # Change to a writable directory
+save_dir = "/tmp/t5-finetuned"  # Use /tmp/, which is writable
 os.makedirs(save_dir, exist_ok=True)  # Ensure the directory exists
-trainer.save_model(save_dir)  
-
+trainer.save_model(save_dir)  # Save the model
 # Load dataset
 dataset = load_dataset("tatsu-lab/alpaca")
 dataset["train"] = dataset["train"].select(range(2000))