Skip to content

Fix CheckpointLoader bug with strict_shape=False and multiple objects #8948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
luomi16 wants to merge 1 commit into
base: main
Choose a base branch
from
+7 −1
Open

Fix CheckpointLoader bug with strict_shape=False and multiple objects #8948

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion monai/handlers/checkpoint_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,14 @@ def __call__(self, engine: Engine) -> None:
checkpoint = {k: checkpoint}

# skip items that don't match data shape
# only apply copy_model_state for torch.nn.Module objects
# this is needed to handle shape mismatches for models
for k, obj in self.load_dict.items():
checkpoint[k] = copy_model_state(obj, checkpoint, inplace=False)[0]
if isinstance(obj, torch.nn.Module) and k in checkpoint:
_, _, _ = copy_model_state(obj, checkpoint[k], inplace=True)
# after copy_model_state, the model's state_dict has been updated
# with matching shapes, so we need to update the checkpoint
checkpoint[k] = obj.state_dict()

# save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint
prior_max_epochs = engine.state.max_epochs
Expand Down