Support AST rewriting for submodules (#52297)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52297
Before, an `nn.Module` with submodules would fail AST rewriting with `TypeError: 'RewrittenModule' object does not support item assignment`. (Try the `test_ast_rewriter_reassigns_submodules` test case on `master`.) This PR fixes the issue as well as adding additional test cases
Test Plan: Imported from OSS
Reviewed By: pbelevich
Differential Revision: D26483820
Pulled By: ansley
fbshipit-source-id: 757e898dc2b0a67daf2bd039d555b85f4e443322
