Reset grad state across unittests (#126345)

Pull Request resolved: #126345 Approved by: https://github.com/ezyang
pytorch · May 23, 2024 · d11e44c · d11e44c
1 parent a31a60d
commit d11e44c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
@@ -102,13 +102,7 @@
 
 
 class AOTTestCase(TestCase):
-    def setUp(self):
-        self.prev_grad_state = torch.is_grad_enabled()
-        super().setUp()
-
-    def tearDown(self):
-        torch.set_grad_enabled(self.prev_grad_state)
-        super().tearDown()
+    pass
 
 
 class TestPythonKey(AOTTestCase):

diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
@@ -2903,6 +2903,9 @@ def setUp(self):
         if self._default_dtype_check_enabled:
             assert torch.get_default_dtype() == torch.float
 
+        # attempt to reset some global state at the end of the test
+        self._prev_grad_state = torch.is_grad_enabled()
+
     def tearDown(self):
         # There exists test cases that override TestCase.setUp
         # definition, so we cannot assume that _check_invariants
@@ -2917,6 +2920,10 @@ def tearDown(self):
         if self._default_dtype_check_enabled:
             assert torch.get_default_dtype() == torch.float
 
+        # attribute may not be defined, per above
+        if hasattr(self, '_prev_grad_state'):
+            torch.set_grad_enabled(self._prev_grad_state)
+
     @staticmethod
     def _make_crow_indices(n_rows, n_cols, nnz,
                            *, device, dtype, random=True):