diff options
Diffstat (limited to 'src/training/run_experiment.py')
| -rw-r--r-- | src/training/run_experiment.py | 7 |
1 files changed, 6 insertions, 1 deletions
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": |