__init__ on nn.Modules runs in Python (torch.jit.script only sees the module after it has been initialized), so you can annotate that with whatever Python annotations you want (but they won’t be enforced by the compiler at all, for that you should use something like mypy).
For methods that are compiled (e.g. forward and anything it calls), the return types can be deduced from the code. If you want to explicitly write it out, you can use any of the type annotations listed here.
Unions aren’t supported so they won’t work in TorchScript. As a workaround you could do something like Tuple[Optional[int], Optional[List[str]]].