From 4a54d7e690897dd6e6c719fb908fd371a44c2952 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Sun, 24 Jan 2021 22:14:17 +0100 Subject: Many updates, cool stuff on the way. --- src/training/run_experiment.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) (limited to 'src/training/run_experiment.py') diff --git a/src/training/run_experiment.py b/src/training/run_experiment.py index 2c9a196..faafea6 100644 --- a/src/training/run_experiment.py +++ b/src/training/run_experiment.py @@ -296,7 +296,12 @@ def run_experiment( # Run inference over test set. if test: logger.info("Loading checkpoint with the best weights.") - model.load_from_checkpoint(model_dir / "best.pt") + if "checkpoint" in experiment_config["train_args"]: + model.load_from_checkpoint( + model_dir / experiment_config["train_args"]["checkpoint"] + ) + else: + model.load_from_checkpoint(model_dir / "best.pt") logger.info("Running inference on test set.") if experiment_config["criterion"]["type"] == "EmbeddingLoss": -- cgit v1.2.3-70-g09d2