Ensure the original grad_mode is restored during backward (#31884)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/31884
Fix #31715
Test Plan: Imported from OSS
Differential Revision: D19301076
Pulled By: albanD
fbshipit-source-id: 2d20c01bfb6364fa96c8fe5aa5ce7ea39defa3ce
