How to add annotation to work with JIT for nn.Sequential

jit
1 
class Model(mm.Module):
    def __init__(self):
        # Init

    def forward(self, x):
        # type: (List[Tensor]) -> List[Tensor]
        # Code

If we need to define type of input x that if different from Tensor, we could do some thing like above. However, i am confused about how to add annotation if i use

class Model(mm.Module):
    def __init__(self):
        # Add sequential module
        self.layer = nn.Sequential([...])

    def forward(self, x):
        # type: (List[Tensor]) -> List[Tensor]
        # Code

If i define an attribute self.layer of type nn.Sequential and this nn.Sequential requires other data type for example List[Tensor], how can i annotate this to work with torch.jit?

2

This is not possible at the moment. What you can do is extend or make your own version of Sequential with the type annotation you need on its forward method.