Added choices to arguments

Author: PrimozGodec

src/image_colorization/main.py

@@ -3,7 +3,12 @@
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Test model')
-    parser.add_argument('--model', type=str, default="reg_full_model", help="Choose name of the model you want to test")
+    parser.add_argument('--model',
+                        type=str,
+                        default="reg_full_model",
+                        help="Choose name of the model you want to test",
+                        choices={"reg_full_model", "reg_full_vgg_model", "reg_part_model",
+                                 "class_weights_model", "class_wo_weights_model"})
 
     args = parser.parse_args()