Fix default dtype restoration after model loader errors (#34440)

Co-authored-by: zhisbug <1654062+zhisbug@users.noreply.github.com> Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
2026-08-11 11:33:04 -07:00
parent c58953d90a
commit 93c1bff1d4
1 changed files with 4 additions and 2 deletions
@@ -25,8 +25,10 @@ def set_default_torch_dtype(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
-    yield
-    torch.set_default_dtype(old_dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)


 def _is_moe_model(model_config: ModelConfig, architectures: list[str]) -> bool: