add epoch option (#933)

2026-09-09 11:32:07 +08:00 · 2020-02-25 09:46:01 -05:00
parent 4d36e76cbc
commit 932770771b
1 changed files with 1 additions and 0 deletions
@@ -250,6 +250,7 @@ class LightningTemplateModel(pl.LightningModule):
        parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)

        # training params (opt)
+        parser.add_argument('--epochs', default=20, type=int)
        parser.add_argument('--optimizer_name', default='adam', type=str)
        parser.add_argument('--batch_size', default=64, type=int)
        return parser