From 569e9ecce0b1e42491705854299fb0118966bf95 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sun, 18 Aug 2024 14:02:41 -0400
Subject: [PATCH] Try to fix TE bug.

---
 comfy/ops.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/comfy/ops.py b/comfy/ops.py
index e11d46f9c..2d2438b81 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -21,6 +21,8 @@ import comfy.model_management
 
 
 def cast_to(weight, dtype=None, device=None, non_blocking=False):
+    if weight.device == device:
+        non_blocking = False
     return weight.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 def cast_to_input(weight, input, non_blocking=False):