From e784ef2fcab8f72059418c176919c476f2bf0dca Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 24 Sep 2025 17:15:03 +0900 Subject: [PATCH] Fix spatial dimension mismatch in concept map creation - Fix 'Input and output must have the same number of spatial dimensions' error - Add proper tensor format handling for F.interpolate - Support 2D, 3D, and 4D tensor formats for interpolation - Ensure correct dimension handling for concept attention tensors - Add dimension squeezing to remove extra batch/channel dimensions --- concept_attention.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/concept_attention.py b/concept_attention.py index 92fc08f..2b7e3d3 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -217,12 +217,31 @@ class ConceptAttention: # Resize to target image size target_h, target_w = image_shape + + # Ensure proper tensor format for interpolation + if len(concept_attention.shape) == 2: # [H, W] + # Add batch and channel dimensions: [1, 1, H, W] + concept_attention_4d = concept_attention.unsqueeze(0).unsqueeze(0) + elif len(concept_attention.shape) == 3: # [B, H, W] + # Add channel dimension: [B, 1, H, W] + concept_attention_4d = concept_attention.unsqueeze(1) + else: + # Already 4D or other format + concept_attention_4d = concept_attention + + # Resize using proper 4D format concept_attention_resized = F.interpolate( - concept_attention.unsqueeze(0).unsqueeze(0), + concept_attention_4d, size=(target_h, target_w), mode='bilinear', align_corners=False - ).squeeze() + ) + + # Remove extra dimensions if needed + if concept_attention_resized.shape[0] == 1: + concept_attention_resized = concept_attention_resized.squeeze(0) + if concept_attention_resized.shape[0] == 1: + concept_attention_resized = concept_attention_resized.squeeze(0) concept_maps[concept] = concept_attention_resized logger.info(f"Created concept map for '{concept}': {concept_attention_resized.shape}")