Make context manager set original state correctly (#126)

allenai · Dec 20, 2024 · 3505660 · 3505660
1 parent 9e0992b
commit 3505660
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/olmo_core/utils.py b/src/olmo_core/utils.py
@@ -643,4 +643,4 @@ def cuda_sync_debug_mode(debug_mode: Union[int, str]):
         yield
     finally:
         if current_mode is not None:
-            torch.cuda.set_sync_debug_mode(debug_mode)
+            torch.cuda.set_sync_debug_mode(current_mode)