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): ...@@ -164,9 +164,11 @@ def manual_cast_forward(target_dtype):
@contextlib.contextmanager @contextlib.contextmanager
def manual_cast(target_dtype): def manual_cast(target_dtype):
applied = False
for module_type in patch_module_list: for module_type in patch_module_list:
if hasattr(module_type, "org_forward"): if hasattr(module_type, "org_forward"):
continue continue
applied = True
org_forward = module_type.forward org_forward = module_type.forward
if module_type == torch.nn.MultiheadAttention and has_xpu(): if module_type == torch.nn.MultiheadAttention and has_xpu():
module_type.forward = manual_cast_forward(torch.float32) module_type.forward = manual_cast_forward(torch.float32)
...@@ -176,6 +178,8 @@ def manual_cast(target_dtype): ...@@ -176,6 +178,8 @@ def manual_cast(target_dtype):
try: try:
yield None yield None
finally: finally:
if not applied:
return
for module_type in patch_module_list: for module_type in patch_module_list:
if hasattr(module_type, "org_forward"): if hasattr(module_type, "org_forward"):
module_type.forward = 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