加载预训练模型，动转静态图报错。

[CPones]

export_model.py

import paddle
from config import get_config
from cait import build_cait as build_model

# config files in ./configs/
config = get_config('./configs/cait_xxs24_224.yaml')
# build model
model = build_model(config)
# load pretrained weights
model_state_dict = paddle.load('./cait_xxs24_224.pdparams')
model.set_state_dict(model_state_dict)
# dynamic to static
input_spec = paddle.static.InputSpec([None, 3, 224, 224], 'float32', 'image')
model = paddle.jit.to_static(model, input_spec=[input_spec])
# save inferenc model
paddle.jit.save(model, "infer/cait_xxs24_224")

TypeError: In transformed code:

File "/home/aistudio/CaiT/cait.py", line 493, in forward
x = self.forward_features(x)
File "/home/aistudio/CaiT/cait.py", line 481, in forward_features
# Self-Attention blocks
File "/home/aistudio/CaiT/cait.py", line 373, in forward
x = self.attn(x)
File "/home/aistudio/CaiT/cait.py", line 228, in forward
    B, H, C = x.shape # H: num_patches
    qkv = self.qkv(x).chunk(3, axis=-1)
    q, k, v = map(self.transpose_multihead, qkv) #[B, num_heads, num_patches, single_head_dim]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

    q = q * self.scale

File "/home/aistudio/CaiT/cait.py", line 220, in transpose_multihead
new_shape = x.shape[:-1] + [self.num_heads, self.dim_head]

TypeError: can only concatenate tuple (not "list") to tuple

[xperzy]

Thanks for the issue.
We are fixing this problem now and will update once finished.