diff --git a/src/corridorkey_mlx/model/decoder.py b/src/corridorkey_mlx/model/decoder.py index f92aff9..fdcdb13 100644 --- a/src/corridorkey_mlx/model/decoder.py +++ b/src/corridorkey_mlx/model/decoder.py @@ -86,8 +86,8 @@ class DecoderHead(nn.Module): x = up(x) projected.append(x) - # Concatenate along channel dim (last dim in NHWC) - fused = mx.concatenate(projected, axis=-1) # (B, H/4, W/4, embed_dim*4) + # Concatenate in c4, c3, c2, c1 order to match Torch decoder + fused = mx.concatenate(projected[::-1], axis=-1) # (B, H/4, W/4, embed_dim*4) fused = self.linear_fuse(fused) fused = self.bn(fused) fused = nn.relu(fused)