Repository: pytorch
Issue: PyTorch 1.4.0 does not support using `Module` or `ModuleList` in attribute annotations in ScriptModule
Labels: oncall: jit,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

The following piece of code:

```python
import torch
from torch import Tensor, zeros
from torch.nn import Module, Linear, ModuleList
from torch.jit import script


class MyModule(torch.nn.Module):

    __constants__ = ('wrapped',)

    wrapped: Module

    def __init__(self, wrapped: Module):
        super().__init__()
        self.wrapped = wrapped

    def forward(self, input: Tensor):
        return self.wrapped(input)


m = script(MyModule(Linear(3, 3)))
```

would produce the following traceback:

```
/opt/local/miniconda3/envs/py3.6/bin/python /Users/ipwx/projects/tensorkit/debug.py
Traceback (most recent call last):
  File "/Users/ipwx/projects/tensorkit/debug.py", line 19, in <module>
    m = script(MyModule(Linear(3, 3)))
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 1255, in script
    return torch.jit._recursive.recursive_script(obj)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 534, in recursive_script
    return create_script_module(nn_module, infer_methods_to_compile(nn_module))
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 293, in create_script_module
    concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 236, in get_or_create_concrete_type
    concrete_type_builder = infer_concrete_type_builder(nn_module)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 115, in infer_concrete_type_builder
    attr_type = infer_type(name, item)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 79, in infer_type
    attr_type = torch.jit.annotations.ann_to_type(class_annotations[name])
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/annotations.py", line 253, in ann_to_type
    raise ValueError("Unknown type annotation: '{}'".format(ann))
ValueError: Unknown type annotation: '<class 'torch.nn.modules.module.Module'>'
```

Also, the following piece of code:

```python
from typing import *

import torch
from torch import Tensor, zeros
from torch.nn import Module, Linear, ModuleList
from torch.jit import script


class MyModule(torch.nn.Module):

    __constants__ = ('wrapped',)

    wrapped: ModuleList

    def __init__(self, wrapped: List[Module]):
        super().__init__()
        self.wrapped = ModuleList(wrapped)

    def forward(self, input: Tensor):
        output = input
        for m in self.wrapped:
            output = m(output)
        return output


m = script(MyModule([Linear(3, 3), Linear(3, 3)]))
```

would produce the following traceback:

```
Traceback (most recent call last):
  File "/Users/ipwx/projects/tensorkit/debug.py", line 24, in <module>
    m = script(MyModule([Linear(3, 3), Linear(3, 3)]))
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 1255, in script
    return torch.jit._recursive.recursive_script(obj)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 534, in recursive_script
    return create_script_module(nn_module, infer_methods_to_compile(nn_module))
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 293, in create_script_module
    concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 236, in get_or_create_concrete_type
    concrete_type_builder = infer_concrete_type_builder(nn_module)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 115, in infer_concrete_type_builder
    attr_type = infer_type(name, item)
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 79, in infer_type
    attr_type = torch.jit.annotations.ann_to_type(class_annotations[name])
  File "/opt/local/miniconda3/envs/py3.6/lib/python3.6/site-packages/torch/jit/annotations.py", line 253, in ann_to_type
    raise ValueError("Unknown type annotation: '{}'".format(ann))
ValueError: Unknown type annotation: '<class 'torch.nn.modules.container.ModuleList'>'
```

Also, I noticed that, if I removed the type annotation for `wrapped`, no error was thrown, but the code of `wrapped.forward` was not inlined, which was the behavior of PyTorch 1.3.1.  Formally, the following piece of code:

```python
import torch
from torch import Tensor, zeros
from torch.nn import Module, Linear, ModuleList
from torch.jit import script


class MyModule(torch.nn.Module):

    __constants__ = ('wrapped',)

    # wrapped: Module

    def __init__(self, wrapped: Module):
        super().__init__()
        self.wrapped = wrapped

    def forward(self, input: Tensor):
        return self.wrapped(input)


m = script(MyModule(Linear(3, 3)))
print(m.code)
```

would generate the following output:

```
def forward(self,
    input: Tensor) -> Tensor:
  return (self.wrapped).forward(input, )
```

but the same code generates the following output with PyTorch 1.3.1:

```
import __torch__
import __torch__.torch.nn.modules.linear
def forward(self,
    input: Tensor) -> Tensor:
  _0 = self.wrapped
  _1 = _0.weight
  _2 = _0.bias
  if torch.eq(torch.dim(input), 2):
    _3 = torch.__isnot__(_2, None)
  else:
    _3 = False
  if _3:
    bias = ops.prim.unchecked_unwrap_optional(_2)
    ret = torch.addmm(bias, input, torch.t(_1), beta=1, alpha=1)
  else:
    output = torch.matmul(input, torch.t(_1))
    if torch.__isnot__(_2, None):
      bias0 = ops.prim.unchecked_unwrap_optional(_2)
      output0 = torch.add_(output, bias0, alpha=1)
    else:
      output0 = output
    ret = output0
  return ret
```

## Expected behavior

I suppose PyTorch should support `Module` and `ModuleList` in its annotations, just as 1.3.1.  Also, the code of wrapped is expected to be inlined, as a performance optimization.

## Environment

 - PyTorch Version (e.g., 1.0): 1.4.0
 - OS (e.g., Linux): Mac OSX 10.15.3
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): none
 - Python version: 3.6
 - CUDA/cuDNN version: none
 - GPU models and configuration: none
 - Any other relevant information: none


cc @suo