Support more types for ScriptModule #771
Description
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.
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)