Support AMP in nn.parallel (#43102)
Summary:
Take care of the state of autocast in `parallel_apply`, so there is no need to decorate model implementations.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43102
Reviewed By: ngimel
Differential Revision: D23294610
Pulled By: mrshenli
fbshipit-source-id: 0fbe0c79de976c88cadf2ceb3f2de99d9342d762
