Skip to content

GitLab

"...git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "429e5449c12fa66ca489702e0a90900af415498c"
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
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
test/test_models.py
View file @ a5035df5
...@@ -29,8 +29,10 @@ SKIP_BIG_MODEL = os.getenv("SKIP_BIG_MODEL", "1") == "1" ...@@ -29,8 +29,10 @@ 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
yield try:
torch.backends.cudnn.allow_tf32 = previous yield
finally:
torch.backends.cudnn.allow_tf32 = previous
def list_model_fns(module): def list_model_fns(module):
......
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