Skip to content

GitLab

Unverified Commit a5035df5 authored by Nikita Shulga's avatar Nikita Shulga Committed by GitHub
Browse files

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

parent f3d6c6be
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
test/test_models.py
View file @ a5035df5
...@@ -29,7 +29,9 @@ SKIP_BIG_MODEL = os.getenv("SKIP_BIG_MODEL", "1") == "1" ...@@ -29,7 +29,9 @@ SKIP_BIG_MODEL = os.getenv("SKIP_BIG_MODEL", "1") == "1"
def disable_tf32(): def disable_tf32():
previous = torch.backends.cudnn.allow_tf32 previous = torch.backends.cudnn.allow_tf32
torch.backends.cudnn.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False
try:
yield yield
finally:
torch.backends.cudnn.allow_tf32 = previous torch.backends.cudnn.allow_tf32 = previous
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment