Skip to content

[BugFix] Skip missing TensorDictParams submodule when loading state_dict with strict=False - #1734

Merged
vmoens merged 1 commit into
pytorch:mainfrom
javierdejesusda:fix/1472-load-state-dict-missing-entries
Jul 9, 2026
Merged

[BugFix] Skip missing TensorDictParams submodule when loading state_dict with strict=False#1734
vmoens merged 1 commit into
pytorch:mainfrom
javierdejesusda:fix/1472-load-state-dict-missing-entries

Conversation

@javierdejesusda

Copy link
Copy Markdown
Contributor

Description

Loading a state_dict that has no entries under a TensorDictParams submodule
prefix crashes. In _load_from_state_dict, data = data.get(prefix) returns
None when the checkpoint contains nothing for that prefix, and the code then
calls self.data.load_state_dict(None) unconditionally, which raises. This is
the normal path for a strict=False load where a submodule is simply absent
from the checkpoint.

This PR guards the load so a missing submodule under strict=False becomes a
no-op, matching the nn.Module contract:

if data is not None:
    self.data.load_state_dict(data)

Motivation and Context

close #1472

A module holding TensorDictParams cannot be partially restored from a
checkpoint that omits those params, even with strict=False. nn.Module
treats missing keys as a no-op in that case, so TensorDictParams should too
instead of raising.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • My change requires a change to the documentation.
  • I have updated the documentation accordingly.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 8, 2026
@github-actions github-actions Bot added bug Something isn't working Test nn and removed bug Something isn't working labels Jul 8, 2026

@vmoens vmoens left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM thanks!

TensorDictParams._load_from_state_dict resolved the module prefix via
TensorDict.get, which returns None when the incoming checkpoint contains
no entries under the prefix. Passing that None to load_state_dict raised
a TypeError.

Delegate the missing-prefix case to nn.Module's loader so non-strict
loads report the absent parameters without raising and strict loads keep
the standard missing-key error. Cover both modes while preserving the
test process's RNG state.
@vmoens
vmoens force-pushed the fix/1472-load-state-dict-missing-entries branch from 61585f9 to ea6ce2e Compare July 9, 2026 17:51
@github-actions github-actions Bot added the bug Something isn't working label Jul 9, 2026
@vmoens
vmoens merged commit 277933a into pytorch:main Jul 9, 2026
66 of 69 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. nn Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] load_state_dict fails if checkpoint lacks entries for TensorDictParams

2 participants