Save a tensorflow model with a transformer layer
HiI 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 ismodel.save_weights('/content/drive/MyDrive/weights/weights')and the second one ismodel.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!