-
Couldn't load subscription status.
- Fork 25.7k
[jit] Add pickler support for Device #30131
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
This PR adds (un)pickling support for `c10::Device`. It also adds `torch.device` as a type annotation for device attributes.
|
Could we reduce forward-compat problems by splitting this into a reader change that we ship first, and a writer change that we ship second? |
|
This doesn't change any existing behavior (someone has to notice this feature and then go and try to serialize a |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good, comments inline.
torch/csrc/jit/unpickler.cpp
Outdated
| stack_.emplace_back(c10::Device(device_string.toStringRef())); | ||
| }); | ||
| stack_.emplace_back(int64_t(globals_.size() - 1)); | ||
| return opcode; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The early returns that have accumulated in this branch of the switch statement should be cleaned up. It is not clear to me why some cases break, others early return, and some push the globals reference number into the stack. Ideally, nothing inside of the switch statement should be returning early. Perhaps put the global logic in its own function for clarity.
Summary: This PR adds (un)pickling support for `c10::Device`. It also adds `torch.device` as a type annotation for device attributes. ](https://our.intern.facebook.com/intern/diff/18664421/) Pull Request resolved: pytorch#30131 Pulled By: driazati Differential Revision: D18664421 fbshipit-source-id: 64378fb42b2d1bbe2bd86259e5ed10f24b5d1e49
This PR adds (un)pickling support for
c10::Device. It also addstorch.deviceas a type annotation for device attributes.Differential Revision: D18664421