From 1b53a71ddeff43bf2af3cc5c30bee6aca0d7c301 Mon Sep 17 00:00:00 2001 From: Andrea Pozzetti Date: Tue, 10 Feb 2026 03:13:49 +0100 Subject: [PATCH] fixed tensor location mismatch --- nodes/utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nodes/utils.py b/nodes/utils.py index 966815d..8e1a47b 100644 --- a/nodes/utils.py +++ b/nodes/utils.py @@ -87,7 +87,7 @@ def masks_to_comfy_mask(masks): masks: torch.Tensor [N, H, W] or [N, 1, H, W] binary masks Returns: - ComfyUI mask tensor [N, H, W] in range [0, 1] + ComfyUI mask tensor [N, H, W] in range [0, 1] on CPU """ if isinstance(masks, torch.Tensor): # Ensure float type and range [0, 1] @@ -99,7 +99,8 @@ def masks_to_comfy_mask(masks): if masks.ndim == 4 and masks.shape[1] == 1: masks = masks.squeeze(1) - return masks + # Move to CPU to ensure compatibility with downstream nodes + return masks.cpu() elif isinstance(masks, np.ndarray): masks = torch.from_numpy(masks).float() if masks.max() > 1.0: @@ -109,6 +110,7 @@ def masks_to_comfy_mask(masks): if masks.ndim == 4 and masks.shape[1] == 1: masks = masks.squeeze(1) + # Already on CPU since from numpy return masks return masks