问题遇到的现象和发生背景

我在跑一个seq2seq模型，其中编码和解码器用到了CuDNNLSTM，我已经训练好模型，并保存了权重weights文件，但是现在要用cpu跑测试，CuDNNLSTM没办法用了。我尝试把他改成LSTM报错
环境 tensorflow1.8 keras 2.1.6

问题相关代码，请勿粘贴截图

模型代码：

encoder，双层双向LSTM

x = LayerNormalization()(x)
x = OurBidirectional(CuDNNLSTM(z_dim // 2, return_sequences=True))([x, x_mask])
x = LayerNormalization()(x)
x = OurBidirectional(CuDNNLSTM(z_dim // 2, return_sequences=True))([x, x_mask])
x_max = Lambda(seq_maxpool)([x, x_mask])

decoder，双层单向LSTM

y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_max])
y = CuDNNLSTM(z_dim, return_sequences=True)(y)
y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_max])
y = CuDNNLSTM(z_dim, return_sequences=True)(y)
y = SelfModulatedLayerNormalization(z_dim // 4)([y, x_max])

运行结果及报错内容

将CuDNNLSTM改成LSTM后报错

Traceback (most recent call last):
File "E:/nlp/my/mytry/TG_BiLSTM/server.py", line 3, in
from generate_title import gen_sent
File "E:\nlp\my\mytry\TG_BiLSTM\generate_title.py", line 311, in
model.load_weights(sys.path[0]+'/best_model.weights')
File "D:\an\envs\tf1.8\lib\site-packages\keras\engine\topology.py", line 2667, in load_weights
f, self.layers, reshape=reshape)
File "D:\an\envs\tf1.8\lib\site-packages\keras\engine\topology.py", line 3412, in load_weights_from_hdf5_group
K.batch_set_value(weight_value_tuples)
File "D:\an\envs\tf1.8\lib\site-packages\keras\backend\tensorflow_backend.py", line 2372, in batch_set_value
assign_op = x.assign(assign_placeholder)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\ops\variables.py", line 615, in assign
return state_ops.assign(self._variable, value, use_locking=use_locking)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\ops\state_ops.py", line 283, in assign
validate_shape=validate_shape)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\ops\gen_state_ops.py", line 63, in assign
use_locking=use_locking, name=name)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\framework\op_def_library.py", line 787, in _apply_op_helper
op_def=op_def)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\framework\ops.py", line 3392, in create_op
op_def=op_def)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\framework\ops.py", line 1734, in init
control_input_ops)
File "D:\an\envs\tf1.8\lib\site-packages\tensorflow\python\framework\ops.py", line 1570, in _create_c_op
raise ValueError(str(e))
ValueError: Dimension 0 in both shapes must be equal, but are 256 and 512. Shapes are [256] and [512]. for 'Assign_5' (op: 'Assign') with input shapes: [256], [512].

我的解答思路和尝试过的方法

我想要达到的结果