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
Show 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,7 +29,9 @@ 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
+    try:
        yield
+    finally:
        torch.backends.cudnn.allow_tf32 = previous