You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
My model code makes extensive use of tensors created during the forward pass. As part of this, I have to specify the device associated with each tensor to avoid conflicts. I want to export my model in TorchScript for deployment elsewhere, but the DeviceDtypeModuleMixin properties do not appear to be compatible with TorchScript.
While this appears to affect the dtype property too, for my use case, using self.device in a LightningModule prevents successful TorchScript export.
Additionally, the device property returns a Union[str, torch.device], so even if it worked as a property in TorchScript, the return value of the function is incompatible with the type expected by the device keyword option (Optional[Device]). There appears to have been some discussion of this in this issue, but I'm not sure how it was resolved.
I am happy to try contributing a solution, but I wanted to discuss here first. A possible solution seems to involve eliminating the @property decorator and save the current state of the device as a plain device attribute on the model object. The actions currently in the device() function body should be performed as the device attribute is being set rather than on the way out.
Expected behavior
The code shown above is successfully saved as TorchScript.
Hi,
I looked into it and first stumbled upon pytorch/pytorch#37883 which claims that properties are indeed supported. I could also find the related code (here) so I am not very sure why it fails right now. My guess is that the linked code would need to be added to the ScriptModule as well.
while looking for workarounds, I also found the possibility to add it as a raw attribute, but we likely would have to guard it with some logic in __getattr__ and __setattr__ (probably 1:1 the logic we have in the property methods right now).
TorchScript is no longer worked on in PyTorch (as pytorch/pytorch#67146 (comment)). So this will probably stay open forever unless someone contributes the patch
🐛 Bug
My model code makes extensive use of tensors created during the forward pass. As part of this, I have to specify the device associated with each tensor to avoid conflicts. I want to export my model in TorchScript for deployment elsewhere, but the DeviceDtypeModuleMixin properties do not appear to be compatible with TorchScript.
While this appears to affect the
dtype
property too, for my use case, usingself.device
in a LightningModule prevents successful TorchScript export.To Reproduce
results in the following error:
I am not sure how this can be fixed without changing the way DeviceDtypeModuleMixin works. It relies on the
@property
decorator, which uses descriptors. TorchScript is not compatible with descriptors (3.3.2.2, 3.3.2.3).Additionally, the
device
property returns aUnion[str, torch.device]
, so even if it worked as a property in TorchScript, the return value of the function is incompatible with the type expected by thedevice
keyword option (Optional[Device]
). There appears to have been some discussion of this in this issue, but I'm not sure how it was resolved.I am happy to try contributing a solution, but I wanted to discuss here first. A possible solution seems to involve eliminating the
@property
decorator and save the current state of the device as a plaindevice
attribute on the model object. The actions currently in thedevice()
function body should be performed as thedevice
attribute is being set rather than on the way out.Expected behavior
The code shown above is successfully saved as TorchScript.
Environment
conda
,pip
, source): pipcc @carmocca @justusschock @awaelchli @Borda @ananthsub @ninginthecloud @jjenniferdai @rohitgr7
The text was updated successfully, but these errors were encountered: