title: Isues with saving and loading tensorflow model which uses hugging face transformer model as its first layer
body: <p>Hi I am having some serious problems saving and loading a tensorflow model which is combination of hugging face transformers + some custom layers to do classfication. I am using the latest Huggingface transformers tensorflow keras version. The idea is to extract features using distilbert and then run the features through CNN to do classification and extraction. I have got everything to work as far as getting the correct classifications.</p>
  
  <p>The problem is in saving the model once trained and then loading the model again.</p>
  
  <p>I am using tensorflow keras and tensorflow version 2.2</p>
  
  <p>Following is the code to design the model, train it, evaluate it and then save and load it</p>
  
  <pre><code>
      bert_config = DistilBertConfig(dropout=0.2, attention_dropout=0.2, output_hidden_states=False)
      bert_config.output_hidden_states = False
      transformer_model = TFDistilBertModel.from_pretrained(DISTIL_BERT, config=bert_config)
  
      input_ids_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='input_token', dtype='int32')
      input_masks_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='masked_token', dtype='int32')
  
      embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in)[0]
      x = tf.keras.layers.Bidirectional(
          tf.keras.layers.LSTM(50, return_sequences=True, dropout=0.1,
                               recurrent_dropout=0, recurrent_activation="sigmoid",
                               unroll=False, use_bias=True, activation="tanh"))(embedding_layer)
      x = tf.keras.layers.GlobalMaxPool1D()(x)
  
      outputs = []
      # lots of code here to define the dense layers to generate the outputs
      # .....
      # .....
  
      model = Model(inputs=[input_ids_in, input_masks_in], outputs=outputs)
      for model_layer in model.layers[:3]:
          logger.info(f"Setting layer {model_layer.name} to not trainable")
          model_layer.trainable = False
      rms_optimizer = RMSprop(learning_rate=0.001)
      model.compile(loss=SigmoidFocalCrossEntropy(), optimizer=rms_optimizer)
  
      # the code to fit the model (which works)
      # then code to evaluate the model (which also works)
  
      # finally saving the model. This too works.
      tf.keras.models.save_model(model, save_url, overwrite=True, include_optimizer=True, save_format="tf")
  </code></pre>
  
  <p>However, when I try to load the saved model using the following</p>
  
  <pre><code>    tf.keras.models.load_model(
              path, custom_objects={"Addons&gt;SigmoidFocalCrossEntropy": SigmoidFocalCrossEntropy})
  </code></pre>
  
  <p>I get the following load error</p>
  
  <pre><code>
  ValueError: The two structures don't have the same nested structure.
  
  First structure: type=TensorSpec str=TensorSpec(shape=(None, 128), dtype=tf.int32, name='inputs')
  
  Second structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}
  
  More specifically: Substructure "type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}" is a sequence, while substructure "type=TensorSpec str=TensorSpec(shape=(None, 128), dtype=tf.int32, name='inputs')" is not
  Entire first structure:
  .
  Entire second structure:
  {'input_ids': .}
  </code></pre>
  
  <p>I believe the issue is because TFDistilBertModel layer can be called using a dictionary input from DistilBertTokenizer.encode() and that happens to be the first layer. So the model compiler on load expects that to be the input signature to the call model. However, the inputs defined to the model are two tensors of shape (None, 128)</p>
  
  <p>So how do I tell the load function or the save function to assume the correct signatures? </p>