Failing to load saved TFBertModel
TF version: 2.2.0-rc1
transformers version: 2.7.0


`import tensorflow as tf`
`import transformers`
`print(tf.__version__)`
`print(transformers.__version__)`
`MAX_LEN = 10`
`model_path = 'saved_model/temp_model'`
`ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`mask = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`token_type_ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`base_model = transformers.TFBertModel.from_pretrained("bert-base-cased"`
                `, output_hidden_states=False)`
`base_output = base_model([ids, mask, token_type_ids])`
`seq_out, _ = base_output[0], base_output[1]`
`base_model.trainable = False`
`model = tf.keras.models.Model(inputs=[ids, mask, token_type_ids], outputs=[seq_out])`
`model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])`
`print(model.summary())`
`model.save(model_path)`
`model = tf.keras.models.load_model(model_path)`

Model load fails with the following error:


Traceback (most recent call last):
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 378, in assert_same_structure
    expand_composites)
TypeError: The two structures don't have the same nested structure.

First structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "temp.py", line 29, in <module>
    model = tf.keras.models.load_model(model_path)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 190, in load_model
    return saved_model_load.load(filepath, compile)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 116, in load
    model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
    export_dir)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 188, in __init__
    super(KerasObjectLoader, self).__init__(*args, **kwargs)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in __init__
    self._load_all()
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 215, in _load_all
    self._finalize_objects()
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 506, in _finalize_objects
    _finalize_saved_model_layers(layers_revived_from_saved_model)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 677, in _finalize_saved_model_layers
    inputs = infer_inputs_from_restored_call_function(call_fn)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 921, in infer_inputs_from_restored_call_function
    spec = nest.map_structure(common_spec, spec, spec2)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 611, in map_structure
    expand_composites=expand_composites)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 385, in assert_same_structure
    % (str(e), str1, str2))
TypeError: The two structures don't have the same nested structure.

First structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list
Entire first structure:
{'input_ids': .}
Entire second structure:
[., ., .]
Facing a similar issue with tf 2.2.0-rc3.
```
def get_model(lr=0.00001):
    inp_bert = tf.keras.layers.Input(shape=(512), dtype="int32")
    bert = TFBertModel.from_pretrained('bert-base-multilingual-cased')(inp_bert)[0]
    doc_encodings = tf.squeeze(bert[:, 0:1, :], axis=1)
    out = tf.keras.layers.Dense(1, activation="sigmoid")(doc_encodings)
    model = tf.keras.Model(inp_bert, out)
    adam = tf.keras.optimizers.Adam(lr=lr)
    model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
    return model
model = get_model()
model.save("model_name",save_format='tf')
model = tf.keras.models.load_model('model_name')
model.summary()
```
Output error is:

```
/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    383                   "Entire first structure:\n%s\n"
    384                   "Entire second structure:\n%s"
--> 385                   % (str(e), str1, str2))
    386 
    387 

ValueError: The two structures don't have the same nested structure.

First structure: type=TensorSpec str=TensorSpec(shape=(None, 512), 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, 512), dtype=tf.int32, name='inputs')" is not
Entire first structure:
.
Entire second structure:
{'input_ids': .}
```<|||||>change  
`base_output = base_model([ids, mask, token_type_ids])`  
to   
`base_output = base_model.bert([ids, mask, token_type_ids])`
should fix
<|||||>> 
> 
> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

Thanks @Souls362 .. solves it.<|||||>This worked for me as well with `TFBertModel`, however, I run into the same issue with `TFXLNetModel`. `TFXLNetModel` doesn't seem to have an equivalent to the `.bert` property/attribute. Does anyone know how to solve this when using `TFXLNetModel`?<|||||>For `TFXLNetModel` this would be the `.transformer` attribute, as you can see [here](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_tf_xlnet.py#L1127)<|||||>@LysandreJik thank you! That works perfectly<|||||>@LysandreJik How about `TFOpenAIGPTLMHeadModel` ? I use `.transformer` attribute, but the output shape become `[None, None, 768]`, while the original output shape of `TFOpenAIGPTLMHeadModel` is `[None, None, 13088]`. How to solve it? Thanks a lot!<|||||>Well the `transformer` attribute is the transformer in itself, which has a hidden size of 768. The LM head model has an additional head which is the embedding matrix of, which has a size of 13088.<|||||>Yes, I think so too. So how can I save the whole model?<|||||>@Souls362  you are the greatest! I looked way too long for this. 
@huggingface folks, please add an extra detailed example on serialization in TF2. 
There is seriously some clear documentation missing there.<|||||>> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

one tip for TFBertSequenceClassification:  base_model.bert([ids, mask, token_type_ids])[1]<|||||>> > change
> > `base_output = base_model([ids, mask, token_type_ids])`
> > to
> > `base_output = base_model.bert([ids, mask, token_type_ids])`
> > should fix
> 
> one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]

What is the difference of 0 and 1 in the brackets?<|||||>> > > change
> > > `base_output = base_model([ids, mask, token_type_ids])`
> > > to
> > > `base_output = base_model.bert([ids, mask, token_type_ids])`
> > > should fix
> > 
> > 
> > one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]
> 
> What is the difference of 0 and 1 in the brackets?

[TFBertModel documentation](https://huggingface.co/transformers/model_doc/bert.html#transformers.TFBertModel)

model returns sequence output and pooled output (for classification)<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

best answer ever!<|||||>```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_92500/3480800436.py in <module>
----> 1 model_eval = model.evaluate(
      2     dataset_test,
      3     use_multiprocessing=True,
      4     return_dict=True)

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in evaluate(self, x, y, batch_size, verbose, sample_weight, steps, callbacks, max_queue_size, workers, use_multiprocessing, return_dict)
   1387             with trace.Trace('test', step_num=step, _r=1):
   1388               callbacks.on_test_batch_begin(step)
-> 1389               tmp_logs = self.test_function(iterator)
   1390               if data_handler.should_sync:
   1391                 context.async_wait()

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    826     tracing_count = self.experimental_get_tracing_count()
    827     with trace.Trace(self._name) as tm:
--> 828       result = self._call(*args, **kwds)
    829       compiler = "xla" if self._experimental_compile else "nonXla"
    830       new_tracing_count = self.experimental_get_tracing_count()

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    869       # This is the first call of __call__, so we have to initialize.
    870       initializers = []
--> 871       self._initialize(args, kwds, add_initializers_to=initializers)
    872     finally:
    873       # At this point we know that the initialization is complete (or less

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    723     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    724     self._concrete_stateful_fn = (
--> 725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    726             *args, **kwds))
    727 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   2967       args, kwargs = None, None
   2968     with self._lock:
-> 2969       graph_function, _ = self._maybe_define_function(args, kwargs)
   2970     return graph_function
   2971 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3359 
   3360           self._function_cache.missed.add(call_context_key)
-> 3361           graph_function = self._create_graph_function(args, kwargs)
   3362           self._function_cache.primary[cache_key] = graph_function
   3363 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3194     arg_names = base_arg_names + missing_arg_names
   3195     graph_function = ConcreteFunction(
-> 3196         func_graph_module.func_graph_from_py_func(
   3197             self._name,
   3198             self._python_function,

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    988         _, original_func = tf_decorator.unwrap(python_func)
    989 
--> 990       func_outputs = python_func(*func_args, **func_kwargs)
    991 
    992       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    632             xla_context.Exit()
    633         else:
--> 634           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    635         return out
    636 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    975           except Exception as e:  # pylint:disable=broad-except
    976             if hasattr(e, "ag_error_metadata"):
--> 977               raise e.ag_error_metadata.to_exception(e)
    978             else:
    979               raise

TypeError: in user code:

    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1233 test_function  *
        return step_function(self, iterator)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1224 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1217 run_step  **
        outputs = model.test_step(data)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1188 test_step
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/compile_utils.py:387 update_state
        self.build(y_pred, y_true)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/compile_utils.py:317 build
        self._metrics = nest.map_structure_up_to(y_pred, self._get_metric_objects,
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:1159 map_structure_up_to
        return map_structure_with_tuple_paths_up_to(
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:1241 map_structure_with_tuple_paths_up_to
        assert_shallow_structure(
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:847 assert_shallow_structure
        raise TypeError(_STRUCTURES_HAVE_MISMATCHING_TYPES.format(

    TypeError: The two structures don't have the same sequence type. Input structure has type <class 'tuple'>, while shallow structure has type <class 'dict'>.
```

```python
def map_tk(X, Y):
  X = tokenizer(
    X,
    max_length=max_len,
    padding='max_length',
    truncation=True,
    return_token_type_ids=False,
    return_tensors="tf")
  
  X = {
    "input_ids": tf.reshape(X["input_ids"], [max_len]),
    "attention_mask": tf.reshape(X["attention_mask"], [max_len])
  }
  Y = {
    "y1": to_categorical(Y["y1"], num_classes=11), 
    "y2": to_categorical(Y["y2"], num_classes=4)
    }
  return X, Y


def gen_data(df: pd.DataFrame):
  def gen():
    for _, row in df.iterrows():
      d = {
        "X": row["content"], 
        "Y": {
          "y1": row["y1"],
          "y2": row["y2"]
        }
      }
      yield map_tk(d["X"], d["Y"])
  
  return gen


output_signature = (
  {
    "input_ids": tf.TensorSpec(shape=(150,), dtype=tf.int32),
    "attention_mask": tf.TensorSpec(shape=(150,), dtype=tf.int32)
  },
  {
    "institution": tf.TensorSpec(shape=(11,), dtype=tf.int32),
    "laws_nature": tf.TensorSpec(shape=(4,), dtype=tf.int32)
  })


def build_dataset(df: pd.DataFrame, shuffle_size=0):
  ds = tf.data.Dataset.from_generator(
    gen_data(df),
    output_signature=output_signature)

  if shuffle_size > 0:
    ds = ds.shuffle(buffer_size=shuffle_size)

  return ds.batch(batch_size=batch_size).prefetch(1)


dataset_train = build_dataset(train, 25600)
dataset_valid = build_dataset(valid)
dataset_test = build_dataset(test)
```

This is the inputs of my model, is there any workaround for `electra`?
```python
input_ids = Input(shape=(max_len,), name="input_ids", dtype="int32")
attention_mask = Input(shape=(max_len,), name="attention_mask", dtype="int32")
inputs = {"input_ids": input_ids, "attention_mask": attention_mask}

X = pretrained(inputs)["hidden_states"][-3:-1]
```