From 81126027f5226e7ee58e1a99194eb9ec7b8ec6e7 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sat, 20 Jan 2024 16:31:12 +0800
Subject: [PATCH] Avoid early disable

---
 modules/devices.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modules/devices.py b/modules/devices.py
index 370286297..3bde16990 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -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