Repository: pytorch
Issue: TorchScript: Union cannot be used as a tuple
Labels: oncall: jit

### 🐛 Describe the bug

I can run torch.jit.script on a model, A,  that returns an `Union`, but when I run jit.script on a model, B, that uses A I get an error.

The following code demonstrates the issue:
```python
import torch
from torch import Tensor
from typing import Union, Tuple


class Mymod(torch.nn.Module):

    def __init__(self, return3=False):
        super(Mymod, self).__init__()
        self.return3 = return3

    def forward(
        self, z: Tensor
    ) -> Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]]:
        if self.return3:
            return z, z, z
        else:
            return z, z


class MymodWrapper(torch.nn.Module):
    def __init__(self):
        super(MymodWrapper, self).__init__()
        self.model = Mymod()
        
    def forward(
        self, z: Tensor
    ) -> Tensor:
        o1,_ = self.model(z)
        return o1

mymod = torch.jit.script(Mymod())
o1, _ = mymod(z)
mymodwrapper = torch.jit.script(MymodWrapper()) # Line 60, fails in this line
o1 = mymodwrapper(z)
```  
Fails with the following error:
```python
Traceback (most recent call last):
  File "/root//test.py", line 77, in <module>
    mymodwrapper = torch.jit.script(MymodWrapper())
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/conda/lib/python3.11/site-packages/torch/jit/_script.py", line 1338, in script
    return torch.jit._recursive.create_script_module(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/conda/lib/python3.11/site-packages/torch/jit/_recursive.py", line 558, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/conda/lib/python3.11/site-packages/torch/jit/_recursive.py", line 635, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/conda/lib/python3.11/site-packages/torch/jit/_recursive.py", line 467, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
RuntimeError: 
Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]] cannot be used as a tuple:
  File "/root/test.py", line 60
    def forward(self, z: Tensor) -> Tensor:
    
        o1, b = self.model(z)
                ~~~~~~~~~~~~ <--- HERE
    
        return o1
```

### Versions

```shell
Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.26.4
[pip3] torch==2.3.0.dev20240313
[pip3] torchmetrics==1.3.1
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl    conda-forge
[conda] brotlipy                  0.7.0           py311h9bf148f_1002    pytorch-nightly
[conda] cffi                      1.15.1          py311h9bf148f_3    pytorch-nightly
[conda] cryptography              38.0.4          py311h46ebde7_0    pytorch-nightly
[conda] filelock                  3.9.0                   py311_0    pytorch-nightly
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.2.1         h84fe81f_16997    conda-forge
[conda] mpmath                    1.2.1                   py311_0    pytorch-nightly
[conda] numpy                     1.26.4          py311h64a7726_0    conda-forge
[conda] packaging                 22.0                    py311_0    pytorch-nightly
[conda] pluggy                    1.0.0                   py311_1    pytorch-nightly
[conda] pysocks                   1.7.1                   py311_0    pytorch-nightly
[conda] pytest                    7.1.2                   py311_0    pytorch-nightly
[conda] pytorch                   2.3.0.dev20240313    py3.11_cpu_0    pytorch-nightly
[conda] pytorch-mutex             1.0                         cpu    pytorch-nightly
[conda] requests                  2.28.1                  py311_0    pytorch-nightly
[conda] tomli                     2.0.1                   py311_0    pytorch-nightly
[conda] torchmetrics              1.3.1              pyhd8ed1ab_0    conda-forge
[conda] urllib3                   1.26.14                 py311_0    pytorch-nightly
```
Tried on non nightly pytorch versions too.

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel