Fix handling of type comments in body (#30590)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/30477. Any type comment after `# type: (...) -> ` is ignored.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/30590
Differential Revision: D18887351
Pulled By: driazati
fbshipit-source-id: 162c652f6d7610d14609bbcb25aaa27cdd947a76
