✨ Add a CLI arg to enable resuming training

ENSTA-U2IS-AI · Aug 26, 2023 · 9ed8c4c · 9ed8c4c
1 parent 13a6439
commit 9ed8c4c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/tests/models/test_resnets.py b/tests/models/test_resnets.py
@@ -78,7 +78,7 @@ def test_main(self):
 
 
 class TestMIMOResnet:
-    """Testing the ResNet mimo class."""
+    """Testing the ResNet MIMO class."""
 
     def test_main(self):
         model = mimo_resnet34(1, 10, 2, style="cifar")

diff --git a/torch_uncertainty/__init__.py b/torch_uncertainty/__init__.py
@@ -47,6 +47,11 @@ def init_args(
         action="store_true",
         help="Use channels last memory format",
     )
+    parser.add_argument(
+        "--enable_resume",
+        action="store_true",
+        help="Allow resuming the training (save optimizer's states)",
+    )
 
     parser = pl.Trainer.add_argparse_args(parser)
     if network is not None:
@@ -105,7 +110,7 @@ def cli_main(
         monitor=monitor,
         mode=mode,
         save_last=True,
-        save_weights_only=True,
+        save_weights_only=not args.enable_resume,
     )
 
     # Select the best model, monitor the lr and stop if NaN