Update caliban_toolbox/dataset_splitter.py

Co-authored-by: willgraf <7930703+willgraf@users.noreply.github.com>
vanvalenlab · Sep 30, 2020 · e1bd4f9 · e1bd4f9
1 parent 16bfd86
commit e1bd4f9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/caliban_toolbox/dataset_splitter.py b/caliban_toolbox/dataset_splitter.py
@@ -55,7 +55,7 @@ def __init__(self, seed=0, splits=None):
             self.splits = splits
 
     def _validate_dict(self, train_dict):
-        if 'X' not in train_dict.keys() or 'y' not in train_dict.keys():
+        if 'X' not in train_dict or 'y' not in train_dict:
             raise ValueError('X and y must be keys in the training dictionary')
 
     def split(self, train_dict):