Repository: pytorch
Issue: JIT RuntimeError:  'Union[Tensor, List[float], List[int]]' object is not subscriptable
Labels: oncall: jit

## 🐛 Bug

A properly type hinted function or class results in an error with TorchScript is you slice the sequential List or Tensor object.


The issue only seems to occur when the `List` type is used alongside the `torch.Tensor` type. Otherwise it appears to work normally.

## To Reproduce

Run this test function and convert it to JIT:

```
from typing import Union, List
import torch

def test(x: Union[List[int], List[float], torch.Tensor]) -> Union[int, float, torch.Tensor]:
      return x[0]
out = torch.jit.script(test)
```

Results in the following error message

```
RuntimeError                              Traceback (most recent call last)

<ipython-input-3-c890700ed60b> in <module>()
      4 def test(x: Union[List[int], List[float], torch.Tensor]) -> Union[int, float, torch.Tensor]:
      5       return x[0]
----> 6 out = torch.jit.script(test)

/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1309             _rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
   1310         fn = torch._C._jit_script_compile(
-> 1311             qualified_name, ast, _rcb, get_default_args(obj)
   1312         )
   1313         # Forward docstrings

RuntimeError: 
'Union[Tensor, List[float], List[int]]' object is not subscriptable:
  File "<ipython-input-3-c890700ed60b>", line 5
def test(x: Union[List[int], List[float], torch.Tensor]) -> Union[int, float, torch.Tensor]:
      return x[0]
             ~~~~ <--- HERE
```


## Expected behavior

There should be no error message, unless I'm missing something major.

## Environment


```

Collecting environment information...
PyTorch version: 1.10.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0
Libc version: glibc-2.26

Python version: 3.7.12 (default, Sep 10 2021, 00:21:48)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.10.0+cu111
[pip3] torchaudio==0.10.0+cu111
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.11.0
[pip3] torchvision==0.11.1+cu111
[conda] Could not collect
```

Basically the standard Google Colab environment.

## Additional context

This issue is interfering with making my custom transforms JIT compatible.