Fix silent() not handling exceptions properly

UCSD-PL · Mar 16, 2022 · 16968d5 · 16968d5
1 parent ed258b0
commit 16968d5
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/util.py b/src/util.py
@@ -213,9 +213,11 @@ def silent():
     save_stdout = sys.stdout
     sys.stderr = DummyFile()
     sys.stdout = DummyFile()
-    yield
-    sys.stderr = save_stderr
-    sys.stdout = save_stdout
+    try:
+        yield
+    finally:
+        sys.stderr = save_stderr
+        sys.stdout = save_stdout
 
 with silent():
     use_cuda = torch.cuda.is_available()