-
Notifications
You must be signed in to change notification settings - Fork 3.4k
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
Support torch.jit.script on LightningModules #1951
Comments
The |
i thought we made it compatible in september? |
Does the code block I shared work for you? Are there any tests that actually try to do |
Okay, I'm
and I still see this issue. Perhaps the annotations were added after September and there weren't any tests that caught it? |
if isinstance(dst_type, str):
dst_type = getattr(torch, dst_type) |
I think this issue is fixed after #2657. I am able to convert the model to Torchscript format now. |
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions. |
🚀 Feature
There are a number of advantages to converting a model with TorchScript (e.g. static optimizations, better saving / loading, especially into non-Python environments for deployment). However, no
LightningModule
s can be converted usingtorch.jit.script
. Here's a simple example with the error produced (note that this works as-is if we inherit fromnn.Module
instead ofpl.LightningModule
):Digging into this a little, we have
and the
_device
annotation comes fromDeviceDtypeModuleMixin
(one of the super-classes ofLightningModule
). Here's the relevant snippet:This seems to be the only issue because this code works:
However, if I try to set the annotation to
typing.Union[str, torch.device]
(which seems to be the true type based on this line), then I getValueError: Unknown type annotation: 'typing.Union[str, torch.device]' in
torch.jit.script`.Is the
str
type for_device
actually used? I don't see that anywhere, and I actually do see at least one place where there would be an error ifself.device
returned a string (here. I'll just go ahead and submit a PR to update the annotations, but feel free to comment here or on the PR if there's something I'm missing about the type annotations here.The text was updated successfully, but these errors were encountered: