Repository navigation
Fix loading of palettizer state dict with no indices into prepared model - #147
Merged
crowbat merged 1 commit intoOct 27, 2026
Conversation
NehalBhandari
approved these changes
Oct 27, 2026
crowbat
deleted the
u/k_hsieh/fix_palettizer_loading_sd_centroids_only
branch
October 27, 2026 22:20
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.
In existing code, there is logic to handle loading a KMeansPalettizer state dict with no indices (which would occur when the indices were stale at the time of saving out the state dict).
Existing unit tests only tested loading of this state dict into 1) a bare uninitialized KMeansPalettizer standalone module, or 2) via
prepare()'s dedicated state_dict path which skips module initialization when preparing.However, when loading the state dict into a prepared model with some indices already, I ran into an error due to strict mode flagging the state dict as missing a key. This PR updates the
load_from_state_dictoverride to allow strict loading to succeed. When indices are missing, theindices_staleflag is still marked as True, so new indices would be recomputed at the next forward pass.