Update train.py

train.py

@@ -23,6 +23,7 @@ if "test" not in dataset:
 model_name = "t5-large"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
+model.gradient_checkpointing_enable() 
 
 # Define tokenization function
 def tokenize_function(examples):
@@ -46,14 +47,16 @@ print("Dataset successfully split and tokenized.")
 
 #  Define training arguments
 training_args = TrainingArguments(
-    output_dir=output_dir,
-    per_device_train_batch_size=2,  # Lowered to avoid memory issues
-    per_device_eval_batch_size=2,
+    output_dir="/tmp/t5-finetuned",
+    per_device_train_batch_size=1,  # ✅ Reduce to 1 (was 2)
+    per_device_eval_batch_size=1,  # ✅ Reduce to 1
     num_train_epochs=1,  # Test run (increase for full fine-tuning)
+    gradient_accumulation_steps=4,  # ✅ Helps simulate larger batch size
     logging_steps=50,
     evaluation_strategy="epoch",
     save_strategy="epoch",
-    push_to_hub=False  # Change to True to upload the model to HF Hub
+    push_to_hub=False
+    fp16=True
 )
 
 # Set up Trainer