Repository: pytorch
Issue: Support Slicing of ModuleList during JIT model tracing/scripting 
Labels: oncall: jit,module: bootcamp,triaged,large,medium

## 🐛 Bug

Torch.jit.trace fails if model slicing is used anywhere within the model functions. 

Steps to reproduce the behavior:

Say you have VGG16 model like below - 

```python3
class VGG16_frontend(nn.Module):
    def __init__(self,block_num=5,decode_num=0,load_weights=True,bn=False,IF_freeze_bn=False):
        super(VGG16_frontend,self).__init__()
        self.block_num = block_num
        self.load_weights = load_weights
        self.bn = bn
        self.IF_freeze_bn = IF_freeze_bn
        self.decode_num = decode_num

        block_dict = [[64, 64, 'M'], [128, 128, 'M'], [256, 256, 256, 'M'],\
             [512, 512, 512,'M'], [512, 512, 512,'M']]

        self.frontend_feat = []
        for i in range(block_num):
            self.frontend_feat += block_dict[i]

        if self.bn:
            self.features = make_layers(self.frontend_feat, batch_norm=True)
        else:
            self.features = make_layers(self.frontend_feat, batch_norm=False)


        if self.load_weights:
            if self.bn:
                pretrained_model = models.vgg16_bn(pretrained = True)
            else:
                pretrained_model = models.vgg16(pretrained = True)
            pretrained_dict = pretrained_model.state_dict()
            model_dict = self.state_dict()
            # filter out unnecessary keys
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            # overwrite entries in the existing state dict
            model_dict.update(pretrained_dict) 
            # load the new state dict
            self.load_state_dict(model_dict)

        if IF_freeze_bn:
            self.freeze_bn()
    
    def forward(self,x):
        if self.bn: 
            x = self.features[ 0:7](x)   # <-- Scripting fails here
            conv1_feat =x if self.decode_num>=4 else []
            x = self.features[ 7:14](x)
            conv2_feat =x if self.decode_num>=3 else []
            x = self.features[ 14:24](x)
            conv3_feat =x if self.decode_num>=2 else []
            x = self.features[ 24:34](x)
            conv4_feat =x if self.decode_num>=1 else []
            x = self.features[ 34:44](x)
            conv5_feat =x 
        else:
            x = self.features[ 0: 5](x)
            conv1_feat =x if self.decode_num>=4 else []
            x = self.features[ 5:10](x)
            conv2_feat =x if self.decode_num>=3 else []
            x = self.features[ 10:17](x)
            conv3_feat =x if self.decode_num>=2 else []
            x = self.features[ 17:24](x)
            conv4_feat =x if self.decode_num>=1 else []
            x = self.features[ 24:31](x)
            conv5_feat =x 
               
        feature_map = {'conv1':conv1_feat,'conv2': conv2_feat,\
            'conv3':conv3_feat,'conv4': conv4_feat, 'conv5': conv5_feat}   
        
        # feature_map = [conv1_feat, conv2_feat, conv3_feat, conv4_feat, conv5_feat]
        
        return feature_map


    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()
```
**Command to run**
```python3
example = torch.rand(1, 3, 768, 1024)
fr_end = VGG16_frontend(5,4,True,False,False)
VGG_traced = torch.jit.script(fr_end,example)
```
**Error Message -** 

```python3 
RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  aten::slice.Tensor(Tensor(a) self, int dim=0, int start=0, int end=9223372036854775807, int step=1) -> (Tensor(a)):
  Expected a value of type 'Tensor' for argument 'self' but instead found type '__torch__.torch.nn.modules.container.Sequential'.
  
  aten::slice.str(str string, int start, int end=9223372036854775807, int step=1) -> (str):
  Expected a value of type 'str' for argument 'string' but instead found type '__torch__.torch.nn.modules.container.Sequential'.
  
  aten::slice.t(t[] l, int start, int end=9223372036854775807, int step=1) -> (t[]):
  Could not match type __torch__.torch.nn.modules.container.Sequential to List[t] in argument 'l': Cannot match List[t] to __torch__.torch.nn.modules.container.Sequential.

The original call is:
  File "/home/ubuntu/mayub/Github/SS-DCNet/Network/SSDCNet.py", line 92
    def forward(self,x):
        if self.bn: 
            x = self.features[0:7](x)
                ~~~~~~~~~~~~~~~~~ <--- HERE
            conv1_feat =x if self.decode_num>=4 else []
            x = self.features[ 7:14](x)
```

Expected behavior - should run the trace successfully.

## Environment

```python3
PyTorch version: 1.6.0a0+675e77e
Is debug build: No
CUDA used to build PyTorch: 10.0

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
CMake version: version 3.13.3

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.0.130
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.2

Versions of relevant libraries:
[pip3] numpy==1.15.4
[conda] _tflow_select             2.3.0                       mkl  
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.1.243             h6bb024c_0  
[conda] magma-cuda101             2.5.2                         1    pytorch
[conda] mkl                       2019.4                      243  
[conda] mkl-service               2.3.0            py36he904b0f_0  
[conda] mkl_fft                   1.0.15           py36ha843d7b_0  
[conda] mkl_random                1.1.0            py36hd6b4f25_0  
[conda] numpy                     1.16.2                    <pip>
[conda] numpy                     1.18.3                    <pip>
[conda] numpy                     1.17.4           py36hc1035e2_0  
[conda] numpy-base                1.17.4           py36hde5b4d6_0  
[conda] pytorch                   1.5.0           py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] pytorch-model-summary     0.1.1                     <pip>
[conda] tensorflow                1.15.0          mkl_py36h4920b83_0  
[conda] tensorflow-base           1.15.0          mkl_py36he1670d9_0  
[conda] torch                     1.6.0a0+675e77e           <pip>
[conda] torchfile                 0.1.0                     <pip>
[conda] torchsummary              1.5.1                     <pip>
[conda] torchvision               0.6.0                py36_cu101    pytorch
(Crowd_Detection_mayub) ubuntu@ip-172-31-15-199:~/mayub/Github/example-app/build$ 
```

@suo 

Thanks !


cc @suo