[jit] Fix type annotations in select assignments (#40528)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40528
Previously, an assignment like `self.foo : List[int] = []` would ignore
the type hint.
Test Plan: Imported from OSS
Reviewed By: eellison
Differential Revision: D22222927
Pulled By: suo
fbshipit-source-id: b0af19b87c6fbe0670d06b55f2002a783d00549d
