Make `disable_tf32` context manager handle exceptions (#7651)

a5035df5 · Nikita Shulga · GitHub · f3d6c6be · a5035df5
Unverified Commit a5035df5 authored Jun 02, 2023 by Nikita Shulga Committed by GitHub Jun 02, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

test/test_models.py test/test_models.py +4 -2

No files found.
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -29,8 +29,10 @@ SKIP_BIG_MODEL = os.getenv("SKIP_BIG_MODEL", "1") == "1"
 def disable_tf32():
    previous = torch.backends.cudnn.allow_tf32
    torch.backends.cudnn.allow_tf32 = False
-    yield
-    torch.backends.cudnn.allow_tf32 = previous
+    try:
+        yield
+    finally:
+        torch.backends.cudnn.allow_tf32 = previous


 def list_model_fns(module):