Loading submodules of PyroModule from param store

I have a PyroModule with some normal nn.Module attributes. According to the docs these modules will be synced with the param store. However when I load the param store from a file their parameters are not updated (in the module() call in PyroModule’s __getattr__ method, it doesn’t set update_module_params). Is there an easy way to update these submodules from the param store? Or alternatively should I just be torch.save() and torch.load()'ing the whole PyroModule? The PyroModule docs seem to imply that both methods should work.

@rmccampbell Could you make a simple reproducible script? I can follow up after that. :slight_smile: