fix bug in unimol_tools when using split parameter (#161)

fix bug
deepmodeling · Oct 8, 2023 · 00f67e4 · 00f67e4
1 parent 4ce860f
commit 00f67e4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/unimol_tools/unimol_tools/train.py b/unimol_tools/unimol_tools/train.py
@@ -76,9 +76,9 @@ def update_and_save_config(self):
             self.config['multiclass_cnt'] = self.data['multiclass_cnt']
 
         if self.config['split'] == 'random':
-            self.config['split'] = 'random_5fold'
+            self.config['split_method'] = '5fold_random'
         else:
-            self.config['split'] = 'scaffold_5fold'
+            self.config['split_method'] = '5fold_scaffold'
 
         if self.save_path is not None:
             if not os.path.exists(self.save_path):