fix saving the param store with PyTorch 2 (closes #3201) #3206
Conversation
|
Thanks for fixing this! Hmm, I'm concerned that if we delete the WDYT of this slight modification of your fix # Shallow copy parameters to remove .unconstrained() attributes.
params = {name: param[...] for name, param in self._params.items()}EDIT yup it looks like this is what's causing the failed test. |
In ParamStoreDict.__getitem__, we add a weakref to the unconstrained value to the transformed parameter. Apparently, when the constrained is Real(), transform_to returns the same object that it was passed, so we end up with weakrefs in ParamStoreDict._params. This was fine with PyTorch1, but under Pytorch2 an exception is raised when saving a Parameter with a weakref. So let's explicitly remove the .unconstrained attributes from all parameters to be saved
|
Thanks for the suggestion, fixed. I don't have an overview of the Pyro codebase yet, so I wasn't aware that other Pyro internals rely on |
|
Hmm, looks like some tests are still failing. Let me look into it... |
|
Hey Thanks for your help fixing this @ilia-kats. Following your lead I found more weakref issues and fixed them all at once in #3212. I'll close this PR, but it #3212 wouldn't have happened without your help! |
In ParamStoreDict.getitem, we add a weakref to the unconstrained value to the transformed parameter. Apparently, when the constrained is Real(), transform_to returns the same object that it was passed, so we end up with weakrefs in ParamStoreDict._params. This was fine with PyTorch1, but under Pytorch2 an exception is raised when saving a Parameter with a weakref. So let's explicitly remove the .unconstrained attribute from all parameters to be saved