Testing best model after hyperparameter_search

I am trying to get my head around how to use Ray Tune’s PB2 scheduler alongside the Trainer. Specifically, how can you load the best/final model of a PBT-based scheduler and then continue to test/predict?

With grid search, we could do something like this:

best_params = trainer.hyperparameter_search(...)

# Set the trainer to the best hyperparameters found
for hparam, v in best_params.hyperparameters.items():
    setattr(trainer.args, hparam, v)

# Save optimal hparams
with output_dir.joinpath("opt_hparams.json").open("w", encoding="utf-8") as hp_out:
    dump(best_params, hp_out, indent=4, sort_keys=True)

# Now train the model from-scratch with the best hparams
train_result = trainer.train()
# ... and then get predictions from the model 
predictions = trainer.predict()

but because PB2 changes its hyperparameters on-the-fly, this won’t work because you cannot just do one training run with a single set of parameters.

So how do we typically go about this scenario? How do we load the final, best model after a PB-based hyperparameter-search, and then predict on the test set?

PS I asked this on the Ray forums but did not get a response unfortunately.