Fix CheckpointLoader bug with strict_shape=False and multiple objects

[luomi16]

Summary

This PR fixes a potential bug in CheckpointLoader when using strict_shape=False with multiple objects in load_dict (e.g., both model and optimizer).

Problem

When strict_shape=False and load_dict contains multiple objects (like {'model': model, 'optimizer': optimizer}), the copy_model_state function was being called on all objects, including non-torch.nn.Module objects like optimizers. This caused errors because:

1. copy_model_state expects the source to be iterable (like a state_dict), but optimizers are not iterable
2. Even for models, the state_dict from the checkpoint was never loaded because the key matching logic was incorrect

Solution

The fix ensures that:

1. copy_model_state is only applied to torch.nn.Module objects
2. After copy_model_state updates the model, the checkpoint is updated with the model's new state_dict
3. Other objects (like optimizers) are loaded directly by Checkpoint.load_objects without going through copy_model_state

Testing

• All existing tests pass, including test_strict_shape which specifically tests the strict_shape=False behavior
• The fix has been verified with the test case provided in the issue

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 06c9bce0-383b-4dcc-b530-c927b82679ca

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}