revert one declare_same_as usage

rwth-i6 · Feb 14, 2023 · ffc5083 · ffc5083
1 parent 0cf2cd5
commit ffc5083
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/nn/attention.py b/nn/attention.py
@@ -387,11 +387,10 @@ def __call__(self, spatial_dim: nn.Dim) -> Tuple[nn.Tensor, nn.Dim]:
             remaining_dim = spatial_dim - mat_spatial_size
             left = nn.expand_dim(left, dim=remaining_dim)
             right = nn.expand_dim(right, dim=remaining_dim)
-            concat, out_spatial_dim_ = nn.concat(
+            cond.true, out_spatial_dim_ = nn.concat(
                 (left, remaining_dim), (self.pos_emb, self.clipped_spatial_dim), (right, remaining_dim)
             )
-            concat, out_spatial_dim_ = nn.replace_dim(concat, in_dim=out_spatial_dim_, out_dim=out_spatial_dim)
-            cond.true = concat
+            out_spatial_dim_.declare_same_as(out_spatial_dim)
 
             # False branch, spatial_dim <= self.clipping
             cond.false, _ = nn.slice_nd(