Commit c1713bfe authored by AUTOMATIC1111's avatar AUTOMATIC1111 Committed by GitHub

Merge pull request #14689 from AUTOMATIC1111/fix-nested-manual-cast

Fix nested manual cast
parents 1dbee391 4a66d2fb
......@@ -164,7 +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)
......@@ -174,8 +178,11 @@ def manual_cast(target_dtype):
try:
yield None
finally:
for module_type in patch_module_list:
module_type.forward = module_type.org_forward
if applied:
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward
delattr(module_type, "org_forward")
def autocast(disable=False):
......
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