Save a tensorflow model with a transformer layer

Beginners
1

Hi

I trained a model with the following architecture:

bert_config = BertConfig.from_pretrained(MODEL_NAME)
bert_config.output_hidden_states = True
backbone = TFAutoModelForSequenceClassification.from_pretrained(MODEL_NAME,config=bert_config)

input_ids = tf.keras.layers.Input(shape=(MAX_LENGTH,), name='input_ids', dtype='int32')
features = backbone(input_ids)[1][-1]
pooling =  tf.keras.layers.GlobalAveragePooling1D()(features)
dense = tf.keras.layers.Dense(len(label2id), name='output',activation=tf.nn.softmax)(pooling)
    
model = tf.keras.Model(inputs=[input_ids], outputs = [dense])

and saved model in different ways. The first one is

model.save_weights('/content/drive/MyDrive/weights/weights')

and the second one is

model.save_weights('/content/drive/MyDrive/weights/weights.h5')

So, I am able to load my model (model.load_weights()) from both of these options without any error. Moreover, inference is fine. In short, everything works as I expect.

But if I start a new session and load my model again then inference is bad, like model has random weights instead of my saved weights.

I was trying other options of saving models as well, but they do not work also. Probably there is a special way to save a model with transformer layer?

Thanks in advance!