From 8e5f33cc9c254959826497580819239857c97fef Mon Sep 17 00:00:00 2001 From: drhead <1313496+drhead@users.noreply.github.com> Date: Sun, 9 Mar 2025 19:45:29 -0400 Subject: [PATCH] better way of gating non blocking use --- comfy/conds.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/comfy/conds.py b/comfy/conds.py index 775785d37..e643d6830 100644 --- a/comfy/conds.py +++ b/comfy/conds.py @@ -18,7 +18,10 @@ class CONDRegular: return cond def process_cond(self, batch_size, device, **kwargs): - return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=device_should_use_non_blocking(device)) + if device_should_use_non_blocking(device): + return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True) + else: + return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device) def can_concat(self, other): if self.cond.shape != other.cond.shape: