Support for non Tensor types for `ScriptModule`
- Dominant language
- C#
- Stars
- 1.9k
- Forks
- 228
- PR merge metrics
- No merged PRs in 30d
Description
#771 added some of types but not all. This GH issue is minor.
HuggingFace has `scriptmodel` option in which the returned model only accepts ordinal arguments and outputs a tuple instead of keyword arguments and output `dict`.
```
model = AutoModelForSequenceClassification.from_pretrained(\"distilbert-base-uncased-finetuned-sst-2-english\", torchscript=True)
```
PyTorch mobile's Java binding exports a wrapper of `IValue` (https://pytorch.org/javadoc/1.9.0/)
The Python script below makes a TorchScript file that uses more type for testing.
```python
import torch
from torch import nn
# Used types:
# bool
# bool_list
# dict_long_key
# dict_string_key
# double
# double_list
# list
# long
# long_list
# null
# string
# tensor
# tensor_list
# tuple
class MyModule(nn.Module):
@torch.jit.export
def forward(
self,
x: torch.Tensor,
d: float,
n: int,
s: str,
):
return (
True,
{
1: [False, True],
2: 1.2,
n: [2.3, d],
},
{
"abc": [123, torch.arange(n), None],
s: [456, 789]
},
[x, torch.arange(n)]
)
module = torch.jit.script(MyModule())
print(module.forward.code)
module.save("ivalue_test.pt")
module = torch.jit.load("ivalue_test.pt")
print(module.code)
```
Contributor guide
Assessment
This issue has not been assessed yet.