Machine learning keras-bert迁移学习中的Nan损失

当使用keras bert进行多类分类时，我的损失越来越大。我的输入没有任何NAN值，我也使用了不同的学习速率。我的模型

def build_model(max_seq_length):
   embedding_size = 768
    in_id = tf.keras.layers.Input(shape=(max_seq_length,), name="input_ids")
    in_mask = tf.keras.layers.Input(shape=(max_seq_length,), name="input_masks")
    in_segment = tf.keras.layers.Input(shape=(max_seq_length,), name="segment_ids")
    bert_inputs = [in_id, in_mask, in_segment]

    bert_output = BertLayer(n_fine_tune_layers=3)(bert_inputs)
    bert_output_ = tf.keras.layers.Reshape((max_seq_length, embedding_size))(bert_output)

    LSTM_out = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(100, dropout=0.2))(bert_output_)
    dense = tf.keras.layers.Dense(256, activation="relu")(LSTM_out)
    pred = tf.keras.layers.Dense(5, activation="softmax")(dense)

    model = tf.keras.models.Model(inputs=bert_inputs, outputs=pred)
    adam=tf.keras.optimizers.Adam(learning_rate=0.0001)
    model.compile(loss="sparse_categorical_crossentropy", optimizer=adam, metrics=["accuracy"])
    model.summary()

return model