support multiple split keywords for train and test split

senwu · Nov 25, 2019 · 63e4413 · 63e4413
1 parent 3a04dd9
commit 63e4413
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/emmental/utils/parse_arg.py b/src/emmental/utils/parse_arg.py
@@ -102,7 +102,11 @@ def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser:
     )
 
     learner_config.add_argument(
-        "--train_split", type=str, default="train", help="The split for training"
+        "--train_split",
+        nargs="+",
+        type=str,
+        default=["train"],
+        help="The split for training",
     )
 
     learner_config.add_argument(
@@ -114,7 +118,11 @@ def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser:
     )
 
     learner_config.add_argument(
-        "--test_split", type=str, default="test", help="The split for testing"
+        "--test_split",
+        nargs="+",
+        type=str,
+        default=["test"],
+        help="The split for testing",
     )
 
     learner_config.add_argument(