cant build a sequential model of lstm cells

Question

I had used a code for bayesian lstm by using edward2 library from a paper Bayesian Layers: A Module for Neural Network Uncertainty :

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)

I want to use this code to train a neural network. Could someone help me?