From 1040d02ea0385b09a80ae81b622ae9bb8df4c2f4 Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 17 Sep 2025 14:10:34 +0900 Subject: [PATCH] Fix tensor dimension mismatch in segmentation - Fix RuntimeError: output with shape [1024, 1024] doesn't match broadcast shape [1, 1024, 1024] - Add dimension check and squeeze operation in _create_colored_segmentation - Ensure mask is 2D before broadcasting with colored_image - This fixes the tensor broadcasting issue in ConceptSegmentationNode --- nodes.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nodes.py b/nodes.py index 494c095..240836d 100644 --- a/nodes.py +++ b/nodes.py @@ -291,6 +291,10 @@ class ConceptSegmentationNode: color = colors[concept] mask = (segmentation_mask == i + 1).float() + # Ensure mask is 2D for broadcasting + if len(mask.shape) == 3: + mask = mask.squeeze(0) # Remove batch dimension + for c in range(3): colored_image[0, :, :, c] += mask * color[c]