[internal/CI][PyT/ConvNets] Move to launcher

2021-07-26 21:18:30 +00:00 · 2021-07-26 21:18:30 +00:00 · ab6267ad33
parent d788e8d496
commit ab6267ad33
1 changed files with 7 additions and 0 deletions
--- a/PyTorch/Classification/ConvNets/main.py
+++ b/PyTorch/Classification/ConvNets/main.py
@ -315,6 +315,13 @@ def add_parser_arguments(parser, skip_arch=False):
        choices=[None, "autoaugment"],
        help="augmentation method",
    )
+    parser.add_argument(
+        "--num-classes",
+        type=int,
+        default=None,
+        required=False,
+        help="number of classes"
+    )


 def prepare_for_training(args, model_args, model_arch):