Implementing a callback that loads a checkpoint instead of using ckpt_path="..." in Trainer's fit(), validate(), and test() methods

[ibro45]

Hi all,

I would like to implement a callback that does the same thing that ckpt_path="..." in fit(), validate(), and test() does. Could you please give me a few pointers on what I should do/be careful about?

Thank you!

[ibro45]

https://github.com/Lightning-AI/pytorch-lightning/blob/520c1e4713340f5bbf66de215471c2863e8fbdf2/src/lightning/pytorch/trainer/trainer.py#L954C53-L954C53

Would calling trainer._checkpoint_connector._restore_modules_and_callbacks(ckpt_path) suffice? Any advice on how to deal with trainer.strategy.restore_checkpoint_after_setup in case of such callback implementation?