reduce params to avoid cuda out of memory
@@ -160,8 +160,8 @@ class HFTransformersModel(BaseModel):
"output_dir": self.path.joinpath(self.OUTPUT_DIR_NAME),
"overwrite_output_dir": True,
"num_train_epochs": "30",
- "per_device_train_batch_size": "16",
- "per_device_eval_batch_size": "16",
+ "per_device_train_batch_size": "8",
+ "per_device_eval_batch_size": "8",
"gradient_accumulation_steps": "2",
"learning_rate": "1e-4",
"weight_decay": "0.005",
