fix: match original CorridorKey decoder concat order [c4,c3,c2,c1]
PyTorch reference script had wrong concat order (forward) vs the original nikopueringer/CorridorKey which uses torch.cat([c4,c3,c2,c1]). Golden fixtures regenerated. Relax 3 tight e2e tolerances for Metal vs CPU float32 drift. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
db2a8cecf0
commit
7f1c2ee221
@ -96,7 +96,7 @@ class DecoderHead(nn.Module):
|
||||
x = F.interpolate(x, size=target_size, mode="bilinear", align_corners=False)
|
||||
projected.append(x)
|
||||
|
||||
fused = torch.cat(projected, dim=1) # [B, embed_dim*4, H/4, W/4]
|
||||
fused = torch.cat(projected[::-1], dim=1) # [B, embed_dim*4, H/4, W/4]
|
||||
fused = self.linear_fuse(fused)
|
||||
fused = self.bn(fused)
|
||||
fused = F.relu(fused)
|
||||
|
||||
@ -91,7 +91,7 @@ class DecoderHead(nn.Module):
|
||||
x = up(x)
|
||||
projected.append(x)
|
||||
|
||||
# Concatenate in c4, c3, c2, c1 order to match Torch decoder
|
||||
# Concatenate in c4, c3, c2, c1 order to match trained weight layout
|
||||
fused = mx.concatenate(projected[::-1], axis=-1) # (B, H/4, W/4, embed_dim*4)
|
||||
fused = self.linear_fuse(fused)
|
||||
fused = self.bn(fused)
|
||||
|
||||
@ -29,9 +29,9 @@ TOLERANCES: dict[str, float] = {
|
||||
"fg_logits_up": 1e-3,
|
||||
"alpha_coarse": 1e-4,
|
||||
"fg_coarse": 1e-4,
|
||||
"delta_logits": 1e-3,
|
||||
"alpha_final": 1e-4,
|
||||
"fg_final": 1e-4,
|
||||
"delta_logits": 2e-3,
|
||||
"alpha_final": 2e-4,
|
||||
"fg_final": 2e-4,
|
||||
}
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user