From 3505660866ca46d4341d2ba04306c5450b08bdad Mon Sep 17 00:00:00 2001 From: Shane A Date: Fri, 20 Dec 2024 13:57:49 -0800 Subject: [PATCH] Make context manager set original state correctly (#126) --- src/olmo_core/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/olmo_core/utils.py b/src/olmo_core/utils.py index 64ba1193..d618e827 100644 --- 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)