diff --git a/src/openpi/shared/array_typing.py b/src/openpi/shared/array_typing.py
index 569eafef14..2fa1fc02cd 100644
--- a/src/openpi/shared/array_typing.py
+++ b/src/openpi/shared/array_typing.py
@@ -57,8 +57,10 @@ def typecheck(t: T) -> T:
 def disable_typechecking():
     initial = config.jaxtyping_disable
     config.update("jaxtyping_disable", True)  # noqa: FBT003
-    yield
-    config.update("jaxtyping_disable", initial)
+    try:
+        yield
+    finally:
+        config.update("jaxtyping_disable", initial)
 
 
 def check_pytree_equality(*, expected: PyTree, got: PyTree, check_shapes: bool = False, check_dtypes: bool = False):
diff --git a/src/openpi/shared/array_typing_test.py b/src/openpi/shared/array_typing_test.py
new file mode 100644
index 0000000000..7b54648468
--- /dev/null
+++ b/src/openpi/shared/array_typing_test.py
@@ -0,0 +1,13 @@
+from jaxtyping import config
+import pytest
+
+from openpi.shared import array_typing as at
+
+
+def test_disable_typechecking_restores_flag_on_exception():
+    initial = config.jaxtyping_disable
+
+    with pytest.raises(RuntimeError, match="boom"), at.disable_typechecking():
+        raise RuntimeError("boom")
+
+    assert config.jaxtyping_disable == initial