neural-amp-modeler

commit e8d3dbede3c532b7b079a9c948abe8a7a3f8664f
parent 44bb9e656148586d6a210b57bea59343baf30212
Author: Steven Atkinson <[email protected]>
Date:   Sun, 24 Apr 2022 21:34:29 -0700

Get best checkpoint before showing results

Diffstat:
M
bin/train/main.py
 | 
4
++++

1 file changed, 4 insertions(+), 0 deletions(-)
diff --git a/bin/train/main.py b/bin/train/main.py
@@ -117,6 +117,10 @@ def main(args):
         val_dataloader,
         **learning_config.get("trainer_fit_kwargs", {}),
     )
+    # Go to best checkpoint
+    model = Model.load_from_checkpoint(
+        trainer.checkpoint_callback.best_model_path, **Model.parse_config(model_config)
+    )
     model.eval()
     plot(
         model,