eellison:

Generally, overloads should be usable and work as you expect when they are contained in other modules but you may run into a little bit of difficulty if you are interacting with them as a top level module. They’re still an internal feature but will be cleaned up and released probably in the next release or the one after.


Thank you very much:)
I put overloads into sub modules and find there is one single graph for one kind of overload as you said:
class MySubModule(nn.Module):
    def __init__(self):
        super().__init__()

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: int) -> int:  # noqa: F811
        pass

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: torch.Tensor) -> torch.Tensor:  # noqa: F811
        pass

    def forward(self, x):  # noqa: F811
        original = x

        if isinstance(original, int):
            return original + 2
        else:
            return original.sum()

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.sub_module = MySubModule()

    def forward(self, x):
        return self.sub_module(x)


model = MyModule()
s = torch.jit.script(model)
torch._C._jit_pass_inline(s.graph)
print(s.graph)

output:
graph(%self : __torch__.MyModule,
      %x.1 : Tensor):
      %2 : __torch__.MySubModule = prim::GetAttr[name="sub_module"](%self)
      %14 : None = prim::Constant()
      %15 : Tensor = aten::sum(%x.1, %14) 
  return (%15)

Since default type is torch.Tensor, the corresponding graph is obtained. So is there any ways to obtain one single torchscript for all overloads for now?