pytorch

a80a438e - correctly set and restore states in te tests (#37210)

correctly set and restore states in te tests (#37210) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/37210 Differential Revision: D21238634 Pulled By: Krovatkin fbshipit-source-id: 6462239753399c10c871baa5d5fdff5465cf2544