Issue saving tf Bert model

Question

I have this below model, when Im trying to saving it, Im getting

IndexError: Exception encountered when calling layer "bert" (type TFBertMainLayer).

list index out of range

My model:

input_ids = tf.keras.layers.Input(shape=(MAX_SEQ_LEN,), name='input_token', dtype='int32') input_masks_ids = tf.keras.layers.Input(shape=(MAX_SEQ_LEN,), name='masked_token', dtype='int32') sequence_output = bert_layer.bert([input_ids, input_masks_ids])["last_hidden_state"] x = tf.keras.layers.Lambda(lambda seq: seq[:, 0, :])(sequence_output) out = tf.keras.layers.Dense(3,activation="softmax")(x)

model = tf.keras.Model(inputs=(input_ids, input_masks_ids), outputs=out) model.save("/content/drive/MyDrive/Berttest",save_format='tf')