Update train.py

train.py

@@ -6,9 +6,7 @@ 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 = "/tmp/t5-finetuned"  # Use /tmp/, which is writable
-os.makedirs(save_dir, exist_ok=True)  # Ensure the directory exists
-trainer.save_model(save_dir)  # Save the model
+
 # Load dataset
 dataset = load_dataset("tatsu-lab/alpaca")
 dataset["train"] = dataset["train"].select(range(2000))
@@ -71,6 +69,10 @@ trainer = Trainer(
     eval_dataset=eval_dataset,
 )
 
+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)  # Save the model
+
 # Start fine-tuning
 trainer.train()