Commit 81126027 authored by Kohaku-Blueleaf's avatar Kohaku-Blueleaf

Avoid early disable

parent 0181c1f7
......@@ -164,9 +164,11 @@ def manual_cast_forward(target_dtype):
@contextlib.contextmanager
def manual_cast(target_dtype):
applied = False
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
continue
applied = True
org_forward = module_type.forward
if module_type == torch.nn.MultiheadAttention and has_xpu():
module_type.forward = manual_cast_forward(torch.float32)
......@@ -176,6 +178,8 @@ def manual_cast(target_dtype):
try:
yield None
finally:
if not applied:
return
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment