[BugFix] Skip missing TensorDictParams submodule when loading state_dict with strict=False - #1734
Merged
vmoens merged 1 commit intoJul 9, 2026
Conversation
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
force-pushed
the
fix/1472-load-state-dict-missing-entries
branch
from
July 9, 2026 17:51
61585f9 to
ea6ce2e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Loading a
state_dictthat has no entries under aTensorDictParamssubmoduleprefix crashes. In
_load_from_state_dict,data = data.get(prefix)returnsNonewhen the checkpoint contains nothing for that prefix, and the code thencalls
self.data.load_state_dict(None)unconditionally, which raises. This isthe normal path for a
strict=Falseload where a submodule is simply absentfrom the checkpoint.
This PR guards the load so a missing submodule under
strict=Falsebecomes ano-op, matching the
nn.Modulecontract:Motivation and Context
close #1472
A module holding
TensorDictParamscannot be partially restored from acheckpoint that omits those params, even with
strict=False.nn.Moduletreats missing keys as a no-op in that case, so
TensorDictParamsshould tooinstead of raising.
Types of changes
Checklist