Repository: pytorch
Issue: About ModuleDict indexing when save model based on "torch.jit.script()" method
Labels: oncall: jit

Beblow is the definition of my network.

```python
import torch
import torch.nn as nn


class TinyNet(nn.Module):
    def __init__(self):
        super(TinyNet, self).__init__()
        self.emb_layer = nn.ModuleDict(
            {
                "x": nn.Embedding(10, 10),
                "y": nn.Embedding(10, 10)
            }
        )

    def forward(self, inputs):
        res = list()
        for i in inputs:
            res.append(self.emb_layer[i](0))

        return torch.cat(res)


if __name__ == "__main__":
    net = TinyNet()
    script_model = torch.jit.script(net)
    net.save("pt_model.pt")
```

Run this code will have error  as below:

```python
"Unable to extract string literal index. ModuleDict indexing is only supported with string literals."
res = list()
for i in inputs:
  	res.append(self.emb_layer[i](0))
  ~~~~~~~~~~~~~~~~~ <--- HERE

return torch.cat(res)
```

And I found that issue #45716 seems to have fixed this problem by adding interface type annotation. So I look at the [test/jit/test_module_containers.py](https://github.com/pytorch/pytorch/pull/45716/files/e8cd939ead8e5a8653aabdb555a8c3bb02ce7a37#diff-b43d9b6330ceacf719141b72082934dd3a2dfd69a3b15706a6622276307cc6e6) and found code as below:

```python
@torch.jit.interface
class ModuleInterface(torch.nn.Module):
    def forward(self, inp: Any) -> Any:
        pass

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, inp: Any) -> Any:
        if isinstance(inp, torch.Tensor):
            return torch.max(inp, dim=0)

        return inp

class Mod(torch.nn.Module):
    __annotations__ = {"d": Dict[str, ModuleInterface]}

    def __init__(self):
        super().__init__()
        self.d = torch.nn.ModuleDict({"module": MyModule()})

    def forward(self, x: torch.Tensor, key: str) -> Any:
        return self.d[key].forward(x)


m = Mod()
torch.jit.script(m)
```

So I changed my code as below:

```python
import torch
import torch.nn as nn
from typing import Any, Dict


@torch.jit.interface
class ModuleInterface(torch.nn.Module):
    def forward(self, inp: Any) -> Any:
        pass


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(10, 10)

    def forward(self, inp: torch.Tensor) -> torch.Tensor:
        return self.emb(inp)


class TinyNet(nn.Module):
    __annotations__ = {"emb_layer": Dict[str, ModuleInterface]}

    def __init__(self):
        super(TinyNet, self).__init__()
        self.emb_layer = nn.ModuleDict(
            {
                "x": MyModule(),
                "y": MyModule()
            }
        )

    def forward(self, inputs):
        res = list()
        for i in inputs:
            res.append(self.emb_layer[i](0))
        return torch.cat(res)


if __name__ == "__main__":
    net = TinyNet()
    script_model = torch.jit.script(net)
    net.save("pt_model.pt")
```

But get the error as "  assert attr_type.type().is_interface_type() AssertionError". So in this case, how can I solve the problem? Thank you.

### Environment:

#### Python 3.9.2

#### OS: MAC OSX BigSur 11.6, chip M1

#### torch：1.9.0.post2