python - 无法建立 LSTM 单元的顺序模型

Question

我通过使用论文贝叶斯层中的 edward2 库使用了贝叶斯 lstm 的代码：神经网络不确定性模块：

lstm=ed.layers.LSTMCellReparameterization(8)
output_layer=tf.keras.layers.Dense(4)
def loss_fn(x,lable,datasetsize):
    state = lstm.get_initial_state(x)
    nll = 0.
    for t in range(x.shape[0]):
        net, state = lstm(x, state)
        logits = output_layer(net)
        nll += tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
                              lable, logits=logits))
    k1 = sum(lstm.losses) / datasetsize
    loss=nll+k1
    returnloss
loss1=loss_fn(b1,Y,2000)

我想用这段代码来训练一个神经网络。有人可以帮助我吗？