Update main.py

Jay-Huy · Dec 2, 2024 · 1247fd2 · 1247fd2
1 parent 136782a
commit 1247fd2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/main.py b/main.py
@@ -9,7 +9,7 @@ def parse_args():
         description='Pytorch implementation of MID')
     parser.add_argument('--config', default='', help='Path to the config file')
     parser.add_argument('--dataset', default='', help='Dataset name')
-    parser.add_argument('--network', choices=['ReUNet', 'ReUNet+++'], help='Unet version')
+    parser.add_argument('--network', choices=['ReUNet', 'ReUNet+++', 'Smaller'], help='Unet version')
     parser.add_argument('--data_dir', default=None, help='Path to the data directory')
     parser.add_argument('--epochs', type=int, default=None, help='Number of epochs')
     return parser.parse_args()