Getting Unknown type annotation error when JIT saving

Your example works for me if I remove the declaration of __main outside of __init__ (nn.Module is not a supported type annotation) and rename it to _main (probably some special handling for identifiers that begin with two __).