NVIDIA
🚀[FEA]: allow renaming state-dict keys in from_checkpoint when loading backward-compatible checkpoints
Is this a new feature, an improvement, or a change to existing functionality?
New Feature
How would you describe the priority of this feature request
Critical (currently preventing usage)
Please provide a clear description of problem you would like to solve.
Problem
As reported by @nbren12 the backward compatibility feature allows to rename arguments passed to the constructor when loading an older checkpoint version with from_checkpoint. However, it does not allow to rename saved state_dict keys, which is a common pattern when refactoring model implementations.
Solution
A nice addition to the backward compatibility feature of from_checkpoint would be to be able to map an old_key in the saved state_dict, to a new key in the state_dict of the current model version. Ideally, we would like the mapping to be version-aware, that is, it can adapt the key mapping between the sourfce and target keys based on the Module version.
Describe any alternatives you have considered
No response
