Update app.py

app.py

@@ -141,17 +141,16 @@ def train_model(dataset):
 
     training_args = TrainingArguments(
         output_dir='./results',
-        eval_strategy="epoch",  # Change to eval_strategy
+        evaluation_strategy="epoch",
         learning_rate=2e-5,
-        per_device_train_batch_size=16,
-        per_device_eval_batch_size=16,
+        per_device_train_batch_size=2,  # Try reducing the batch size
+        per_device_eval_batch_size=2,
+        gradient_accumulation_steps=2,  # Accumulate gradients over 2 steps
         num_train_epochs=3,
         weight_decay=0.01,
-        fp16=True,
-        dataloader_num_workers=4,
-        gradient_accumulation_steps=2,
     )
 
+
     class CustomTrainer(Trainer):
         def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             labels = inputs.pop("labels")