Repository: pytorch
Issue: [TorchScript] Failed to Forward Correct Number of Arguments to Different Functions
Labels: oncall: jit

### 🐛 Describe the bug

We are trying to apply `jit.script` to one of our model inside NVIDIA, but find it's hard to support following functions, where the model need to dispatch correct number of arguments to different sub-modules depending on their types, and this type can be decided statically. A minimal example is as follows:

```python
#!/usr/bin/env bash

import torch
from enum import Enum

class ForwardType(Enum):
    One = 1
    Two = 2
    Three = 3

class NetA(torch.nn.Module):
    forward_type : ForwardType = ForwardType.Two
    def __init__(self):
        super(NetA, self).__init__()
        self.linear = torch.nn.Linear(8, 8)

    def forward(self, x, y):
        return self.linear(x) * y

class NetB(torch.nn.Module):
    forward_type : ForwardType = ForwardType.Three
    def __init__(self):
        super(NetB, self).__init__()
        self.linear = torch.nn.Linear(8, 8)

    def forward(self, x, y, z):
        return self.linear(x) * y + z

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        layers = [NetA(), torch.nn.BatchNorm1d(8), NetB()]
        self.layers = torch.nn.ModuleList(layers)

    def forward(self, x):
        # Assume layer can be any unknown Module with 1 to 3 forward args
        # Say there's user defined NetC, NetD in self.layers.
        for layer in self.layers:
            forward_type = getattr(layer, "forward_type", ForwardType.One)
            if forward_type == ForwardType.Two:
                x = layer(x, x)
            elif forward_type == ForwardType.Three:
                x = layer(x, x, x)
            else:
                x = layer(x)
        return x

model = Model()
x = torch.randn(16, 8)
print(model(x).norm())

scripted = torch.jit.script(model)
print(scripted(x).norm())

print(scripted.graph)
```

When we run above snippet with PyTorch `1.14.0a0+410ce96`, we get following error message:

```
$ python tsv2.py
tensor(13.4283, grad_fn=<NormBackward1>)
Traceback (most recent call last):
  File "tsv2.py", line 52, in <module>
    scripted = torch.jit.script(model)
  File "/usr/local/lib/python3.8/dist-packages/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/lib/python3.8/dist-packages/torch/jit/_recursive.py", line 477, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.8/dist-packages/torch/jit/_recursive.py", line 543, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/usr/local/lib/python3.8/dist-packages/torch/jit/_recursive.py", line 394, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError:

forward(__torch__.NetA self, Tensor x, Tensor y) -> Tensor:
Expected at most 3 arguments but found 4 positional arguments.
:
  File "tsv2.py", line 43
                x = layer(x, x)
            elif forward_type == ForwardType.Three:
                x = layer(x, x, x)
                    ~~~~~ <--- HERE
            else:
                x = layer(x)
```

How to resolve the issue?

### Versions

```
Collecting environment information...
PyTorch version: 1.14.0a0+410ce96
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.24.1
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun 22 2022, 20:18:18)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-99-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-80GB
Nvidia driver version: 515.65.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.7.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.2
[pip3] open-clip-torch==2.7.0
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.14.0a0+410ce96
[pip3] torch-tensorrt==1.3.0a0
[pip3] torchtext==0.13.0a0+fae8e8c
[pip3] torchvision==0.15.0a0
[conda] Could not collect
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel