Expose load_state_dict strict=False

[shijianjian]

🚀 Feature

In contrastive learning, we normally train a representation learning backbone then adding the classifier. Sometimes, I wish to play with different classifiers. It would be best to have strict=False exposed to make the model still load with a user's permission.

model = TransferLearningModel.load_from_checkpoint(ckpt_path, strict=False)

The implementation is intuitive but an extremely useful feature in my case. Thank you!

[hbredin]

This feature would also be very useful for LightningModule instances that add a bunch of layers at setup time (as load_from_checkpoint does not seem to call setup).

See related conversation in PytorchLightning Slack.

I'd be happy to try and contribute this (easy, I guess) change if this seems like a good idea.

[s-rog]

yeah seems pretty straight forward. Add a strict arg with default True to load_from_checkpoint that gets passed into _load_model_state and into model.load_state_dict here

unless there are other concerns?

[hbredin]

I don't foresee any other concern but would like to hear from maintainers before spending some time on such a PR.