From f26615fc0c05b695d76c9ac6b23f3b3c29c70bb0 Mon Sep 17 00:00:00 2001
From: John Smith <yfshi123@163.com>
Date: Tue, 28 Mar 2023 21:47:22 +0800
Subject: [PATCH] fix bug

---
 train_data.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train_data.py b/train_data.py
index 62dee6d..76cb2e8 100644
--- a/train_data.py
+++ b/train_data.py
@@ -84,7 +84,7 @@ class TrainTxt(ATrainData):
             r_b = ''
         return new_rows
 
-    def prepare_data(self, thd=-1):
+    def prepare_data(self, thd=-1, **kwargs):
         if os.path.isdir(self.dataset):
             rows = []
             for filename in os.listdir(self.dataset):
@@ -124,7 +124,7 @@ class TrainSAD(ATrainData):
             "attention_mask": result["attention_mask"][:-1],
         }
 
-    def prepare_data(self) -> None:
+    def prepare_data(self, **kwargs) -> None:
         data = load_dataset("json", data_files=self.dataset)
 
         if self.val_set_size > 0: