Skip to content

(fix) Remove nn.Modules and Callables from module init kwargs to enable serialization when weights_only=True#563

Open
dhpitt wants to merge 25 commits intoneuraloperator:mainfrom
dhpitt:load_state_weights_only
Open

(fix) Remove nn.Modules and Callables from module init kwargs to enable serialization when weights_only=True#563
dhpitt wants to merge 25 commits intoneuraloperator:mainfrom
dhpitt:load_state_weights_only

Conversation

@dhpitt
Copy link
Copy Markdown
Member

@dhpitt dhpitt commented Mar 21, 2025

No description provided.

@dhpitt dhpitt changed the title Load state weights only (fix) Remove nn.Modules and Callables from module init kwargs to enable serialization when weights_only=True Mar 21, 2025
@JeanKossaifi
Copy link
Copy Markdown
Member

Looks good - though the tests fail.

@dhpitt
Copy link
Copy Markdown
Member Author

dhpitt commented Mar 24, 2025

Fixed the issues in the doc and in neuralop.models.UQNO causing tests to fail

Copy link
Copy Markdown
Member

@JeanKossaifi JeanKossaifi left a comment

Choose a reason for hiding this comment

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

Thanks @dhpitt, this fixes serialization issue.
I'd add a function though, e.g. get_nonlinearity instead of indexing the dictionary nonlinearity_modules[non_linearity]. The function can still internally use the dict but we can at least catch the KeyError and print a useful message to user.

@JeanKossaifi
Copy link
Copy Markdown
Member

Looks good to me, thank you @dhpitt !

@dhpitt dhpitt mentioned this pull request Jun 2, 2025
@JeanKossaifi
Copy link
Copy Markdown
Member

Was this ready to merge @dhpitt ?

@dhpitt
Copy link
Copy Markdown
Member Author

dhpitt commented Jul 2, 2025 via email

@JeanKossaifi
Copy link
Copy Markdown
Member

@dhpitt getting back to this, do you remember what still needed to be done here?

@dhpitt
Copy link
Copy Markdown
Member Author

dhpitt commented Sep 9, 2025 via email

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants