diff --git a/assets/compile.png b/assets/compile.png
index ac0fde5..ff1e050 100644
Binary files a/assets/compile.png and b/assets/compile.png differ
diff --git a/nodes.py b/nodes.py
index 78b5200..3aa431e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -542,12 +542,12 @@ def teacache_wanmodel_forward_orig(
         self.previous_modulated_input = modulated_inp
 
         if not should_calc:
-            x += self.previous_residual
+            x += self.previous_residual.to(x.device)
         else:
             ori_x = x.clone()
             for block in self.blocks:
                 x = block(x, **kwargs)
-            self.previous_residual = x - ori_x
+            self.previous_residual = (x - ori_x).to(torch.device('cpu'))
 
         # head
         x = self.head(x, e)
diff --git a/pyproject.toml b/pyproject.toml
index 28bfcd1..c02c8f8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "teacache"
 description = "Unofficial implementation of [ali-vilab/TeaCache](https://github.com/ali-vilab/TeaCache) for ComfyUI"
-version = "1.4.0"
+version = "1.4.1"
 license = {file = "LICENSE"}
 
 [project.urls]