Fix SETR export bug

PaddlePaddle · Jul 13, 2021 · 7df00a6 · 7df00a6
1 parent 2bff95e
commit 7df00a6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paddleseg/models/backbones/vision_transformer.py b/paddleseg/models/backbones/vision_transformer.py
@@ -103,7 +103,8 @@ def __init__(self,
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x):
-        N, C = x.shape[1:]
+        x_shape = paddle.shape(x)
+        N, C = x_shape[1], x_shape[2]
         qkv = self.qkv(x).reshape((-1, N, 3, self.num_heads,
                                    C // self.num_heads)).transpose((2, 0, 3, 1,
                                                                     4))