Update train.py

train.py

@@ -24,7 +24,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() 
+model.gradient_checkpointing_enable(False) 
 
 # Define tokenization function
 def tokenize_function(examples):
@@ -58,6 +58,7 @@ training_args = TrainingArguments(
     evaluation_strategy="epoch",
     save_strategy="epoch",
     push_to_hub=False,
+    fp16=True
 )
 
 # Set up Trainer