title: 'tensorflow.keras.models.load_model' results in "ValueError: The two structures don't have the same nested structure."
body: <p>After I save a model and load it, I get "ValueError: The two structures don't have the same nested structure."</p>
  
  <p>This is the code to reproduce the error</p>
  
  <pre><code>!pip install transformers --q
  %tensorflow_version 2.x
  
  from transformers import TFBertModel, AutoModel
  import tensorflow as tf
  from tensorflow.keras.layers import (Dense,
                                       Dropout)
  
  class CustomModel(tf.keras.Model):
      def train_step(self, data):
          # Unpack the data. Its structure depends on your model and
          # on what you pass to `fit()`.
          x, y = data
  
          batch_label = tf.reshape(y, (tf.size(y)/2, 2), name=None)
  
          rs = tf.ragged.stack(x, axis=0)
          reg = rs.to_tensor()
          batch_input = tf.reshape(reg, (tf.shape(reg)[0]*tf.shape(reg)[1], tf.shape(reg)[2]))
  
          with tf.GradientTape() as tape:
              y_pred = self(batch_input, training=True)  # Forward pass
              # Compute the loss value
              # (the loss function is configured in `compile()`)
              loss = self.compiled_loss(batch_label, y_pred, regularization_losses=self.losses)
  
          # Compute gradients
          trainable_vars = self.trainable_variables
          gradients = tape.gradient(loss, trainable_vars)
          # Update weights
          self.optimizer.apply_gradients(zip(gradients, trainable_vars))
          # Update metrics (includes the metric that tracks the loss)
          self.compiled_metrics.update_state(y, y_pred)
          # Return a dict mapping metric names to current value
          return {m.name: m.result() for m in self.metrics}
  
  def get_model(drop_out):
      sciBert = TFBertModel.from_pretrained('allenai/scibert_scivocab_uncased', from_pt=True)
  
      # posFinal = tf.keras.Input(shape=(512,), dtype=tf.int32, name='inputP') 
      # negFinal = tf.keras.Input(shape=(512,), dtype=tf.int32, name='inputN') 
      allFinal = tf.keras.Input(shape=(None,), dtype=tf.int32, name='inputN') 
  
      '''Should posFinal and negFinal be concatenated, so there's only one call to sciBert'''
      # posBertOut = sciBert(posFinal)
      # negBertOut = sciBert(negFinal)
      allBertOut = sciBert(allFinal)
  
      # embedding of size 768*3
      allPoolConcat = tf.concat([
                      allBertOut[1], #output of ff layer after last hidden state since it seems to be untrained in roberta
                      tf.reduce_mean(allBertOut[0][:, 1:-1], axis=1), # pooled except CLS and SEP
                      # tf.math.reduce_max(allBertOut[0][:, 1:-1], axis=1),
                      allBertOut[0][:, -1] # sep, get from hidden state 
                      ],axis=1) 
  
      allPoolConcat = Dropout(drop_out)(allPoolConcat)
      classifier = tf.keras.layers.Dense(2, activation='tanh', name='qff')
  
      allScores = classifier(allPoolConcat) 
  
      model = CustomModel(inputs=allFinal, outputs=allScores)
      return model
  
  config = {
    'drop_out':0.2
  }
  
  loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
  
  model = get_model(**config)
  model.compile(loss=loss_fn,
              optimizer=tf.keras.optimizers.Adam(learning_rate=3e-5))
  
  model.save('modelSave')
  loadedModel = tf.keras.models.load_model('modelSave')
  </code></pre>
  
  <p>And this is the full error message</p>
  
  <pre><code>---------------------------------------------------------------------------
  ValueError                                Traceback (most recent call last)
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
      377     _pywrap_utils.AssertSameStructure(nest1, nest2, check_types,
  --&gt; 378                                       expand_composites)
      379   except (ValueError, TypeError) as e:
  
  ValueError: 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='inputs/input_ids')}
  
  Second structure: type=TensorSpec str=TensorSpec(shape=(None, None), dtype=tf.int32, name='inputs')
  
  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, None), dtype=tf.int32, name='inputs')" is not
  
  During handling of the above exception, another exception occurred:
  
  ValueError                                Traceback (most recent call last)
  11 frames
  &lt;ipython-input-8-4076a08f1890&gt; in &lt;module&gt;()
  ----&gt; 1 loadedModel = tf.keras.models.load_model('modelSave')
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/save.py in load_model(filepath, custom_objects, compile)
      188     if isinstance(filepath, six.string_types):
      189       loader_impl.parse_saved_model(filepath)
  --&gt; 190       return saved_model_load.load(filepath, compile)
      191 
      192   raise IOError(
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in load(path, compile)
      114   # TODO(kathywu): Add saving/loading of optimizer, compiled losses and metrics.
      115   # TODO(kathywu): Add code to load from objects that contain all endpoints
  --&gt; 116   model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
      117 
      118   # pylint: disable=protected-access
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py in load_internal(export_dir, tags, loader_cls)
      602       loader = loader_cls(object_graph_proto,
      603                           saved_model_proto,
  --&gt; 604                           export_dir)
      605       root = loader.get(0)
      606       if isinstance(loader, Loader):
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in __init__(self, *args, **kwargs)
      186     self._models_to_reconstruct = []
      187 
  --&gt; 188     super(KerasObjectLoader, self).__init__(*args, **kwargs)
      189 
      190     # Now that the node object has been fully loaded, and the checkpoint has
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py in __init__(self, object_graph_proto, saved_model_proto, export_dir)
      121       self._concrete_functions[name] = _WrapperFunction(concrete_function)
      122 
  --&gt; 123     self._load_all()
      124     self._restore_checkpoint()
      125 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_all(self)
      213 
      214     # Finish setting up layers and models. See function docstring for more info.
  --&gt; 215     self._finalize_objects()
      216 
      217   @property
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in _finalize_objects(self)
      504         layers_revived_from_saved_model.append(node)
      505 
  --&gt; 506     _finalize_saved_model_layers(layers_revived_from_saved_model)
      507     _finalize_config_layers(layers_revived_from_config)
      508 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in _finalize_saved_model_layers(layers)
      675       call_fn = _get_keras_attr(layer).call_and_return_conditional_losses
      676       if call_fn.input_signature is None:
  --&gt; 677         inputs = infer_inputs_from_restored_call_function(call_fn)
      678       else:
      679         inputs = call_fn.input_signature[0]
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in infer_inputs_from_restored_call_function(fn)
      919   for concrete in fn.concrete_functions[1:]:
      920     spec2 = concrete.structured_input_signature[0][0]
  --&gt; 921     spec = nest.map_structure(common_spec, spec, spec2)
      922   return spec
      923 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in map_structure(func, *structure, **kwargs)
      609   for other in structure[1:]:
      610     assert_same_structure(structure[0], other, check_types=check_types,
  --&gt; 611                           expand_composites=expand_composites)
      612 
      613   flat_structure = [flatten(s, expand_composites) for s in structure]
  
  /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"
  --&gt; 385                   % (str(e), str1, str2))
      386 
      387 
  
  ValueError: 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='inputs/input_ids')}
  
  Second structure: type=TensorSpec str=TensorSpec(shape=(None, None), dtype=tf.int32, name='inputs')
  
  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, None), dtype=tf.int32, name='inputs')" is not
  Entire first structure:
  {'input_ids': .}
  Entire second structure:
  .
  </code></pre>