Make disable_tf32 context manager handle exceptions (#7651)

pytorch · Jun 2, 2023 · a5035df · a5035df
1 parent f3d6c6b
commit a5035df
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -29,8 +29,10 @@
 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):