Lightning-AI / Lightning-AI/pytorch-lightning
Support `assign` argument in `LightningModule.load_from_checkpoint` method
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Currently, `LightningModule.load_from_checkpoint` only allows to pass `strict` argument to `torch.nn.Module.load_state_dict` method. There is also a new `assign` argument (https://github.com/pytorch/pytorch/pull/102212) that can be used to efficiently load model parameters from a checkpoint (see discussion in https://dev-discuss.pytorch.org/t/state-of-model-creation-initialization-seralization-in-pytorch-core/1240). This is a FR to allow passing of `assign` argument to `load_state_dict` similar to the `strict` argument.
If this is something desired I can open a PR.
### Pitch
Implementation can be similar to the `strict` argument. One consideration is to make it backward-compatible (any suggestions on how to?).
### Alternatives
_No response_
### Additional context
_No response_
cc @lantiga @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at LightningModule.load_from_checkpoint and trace where it calls torch.nn.Module.load_state_dict. Follow the existing strict argument handling and determine how assign can be passed while preserving compatibility with older PyTorch versions. Done means callers can provide assign for checkpoint loading without breaking existing strict-only usage.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100