[diffusion] refactor: add arg to control the precision of dit (#17751)

2026-01-27 23:01:23 +08:00
parent 1d942e4eef
commit 547e2d037e
4 changed files with 12 additions and 35 deletions
@@ -20,8 +20,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 get_param_names_mapping(