checkpoint every module in a different ckpt file
Hi!
I am currently working on a project where I would like to checkpoint my model in separated pieces.
My model has a backbone composed by:

a backbone, which is also composed by 3 modules
several heads, each one being a module
I would like to save one ckpt with the backbone and one ckpt per head. I understand that I should create a custom callback inheriting from ModelCheckpoint and then modifying on_save_checkpoint, I am not really aware of how to do it.
on_save_checkpoint is defined as:


Another solution would be to modify my lightning module to load the ckpt when the training ends as a dict and then save each subpart as a ckpt using torch.save(), but I understand that this solution is much less elegant.
Any suggestions? Thanks in advance!
I'd suggest using checkpoint_io plugin for your use-case.