Save model error: list index out of range after pass input_processing call
Author: roymondliao
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.4.0
- Platform: Linux-4.19.0-14-cloud-amd64-x86_64-with-debian-10.8
- Python version: 3.7.10
- PyTorch version (GPU?): 1.7.0 (True)
- Tensorflow version (GPU?): 2.4.1 (True)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
Maybe @LysandreJik or @jplu 
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using albert:

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## Code for Tensorflow
* albert_zh: https://github.com/brightmart/albert_zh

```python
strategy = tf.distribute.MirroredStrategy()
max_seq_len = 128
cache_folder = '/tmp'
pretrain_model = 'voidful/albert_chinese_tiny'
albert_config = AlbertConfig.from_json_file(albert_zh / 'albert_config' / 'albert_config_tiny.json')

def sms_classifier_model(pretrain_model, config, max_seq_len, cache_folder):
    input_ids = tf.keras.layers.Input(shape=(max_seq_len, ), name='input_ids', dtype=tf.int32)
    input_token_type_ids = tf.keras.layers.Input(shape=(max_seq_len, ), name='token_type_ids', dtype=tf.int32)
    input_attention_mask = tf.keras.layers.Input(shape=(max_seq_len, ), name='attention_mask', dtype=tf.int32)
        
    albert_model = TFAlbertForSequenceClassification.from_pretrained(
        pretrain_model, 
        config=config,
        from_pt=True,
        cache_dir=cache_folder)
        
    x = albert_model([input_ids, input_token_type_ids, input_attention_mask])
    output = tf.keras.activations.softmax(x[0])
    model = tf.keras.models.Model(
        inputs=[input_ids, input_token_type_ids, input_attention_mask],
        outputs={'target': output}, name='sms_classifier')
    return model 

K.clear_session()
albert_config.hidden_act = 'gelu_new'
albert_config.num_labels = 4

with strategy.scope():    
    albert_model = sms_classifier_model(pretrain_model, albert_config, 128, cached_pretarin_model_folder)

with strategy.scope():    
    albert_model.compile(optimizer=tf.keras.optimizers.Adam(), 
                         loss=tf.keras.losses.CategoricalCrossentropy(),
                         metrics=tf.keras.metrics.CategoricalAccuracy())

albert_model.fit(
        x=training_dataset,
        validation_data=validation_dataset,
        steps_per_epoch=200,
        validation_steps=100,
        epochs=2,
        verbose=1,
        use_multiprocessing=True)

albert_model.save('/tmp/albert_model')
```
## Error Message
```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-33-9fbc26d706ab> in <module>
      1 albert_model.save(
----> 2     str(saved_tf_model_folder / f'{run_id}')
      3 )

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in save(self, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
   2000     # pylint: enable=line-too-long
   2001     save.save_model(self, filepath, overwrite, include_optimizer, save_format,
-> 2002                     signatures, options, save_traces)
   2003 
   2004   def save_weights(self,

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py in save_model(model, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
    155   else:
    156     saved_model_save.save(model, filepath, overwrite, include_optimizer,
--> 157                           signatures, options, save_traces)
    158 
    159 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save.py in save(model, filepath, overwrite, include_optimizer, signatures, options, save_traces)
     87     with distribution_strategy_context._get_default_replica_context():  # pylint: disable=protected-access
     88       with utils.keras_option_scope(save_traces):
---> 89         save_lib.save(model, filepath, signatures, options)
     90 
     91   if not include_optimizer:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py in save(obj, export_dir, signatures, options)
   1031 
   1032   _, exported_graph, object_saver, asset_info = _build_meta_graph(
-> 1033       obj, signatures, options, meta_graph_def)
   1034   saved_model.saved_model_schema_version = constants.SAVED_MODEL_SCHEMA_VERSION
   1035 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph(obj, signatures, options, meta_graph_def)
   1196 
   1197   with save_context.save_context(options):
-> 1198     return _build_meta_graph_impl(obj, signatures, options, meta_graph_def)

/opt/conda/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph_impl(obj, signatures, options, meta_graph_def)
   1131   if signatures is None:
   1132     signatures = signature_serialization.find_function_to_export(
-> 1133         checkpoint_graph_view)
   1134 
   1135   signatures, wrapped_functions = (

/opt/conda/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py in find_function_to_export(saveable_view)
     73   # If the user did not specify signatures, check the root object for a function
     74   # that can be made into a signature.
---> 75   functions = saveable_view.list_functions(saveable_view.root)
     76   signature = functions.get(DEFAULT_SIGNATURE_ATTR, None)
     77   if signature is not None:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py in list_functions(self, obj, extra_functions)
    149     if obj_functions is None:
    150       obj_functions = obj._list_functions_for_serialization(  # pylint: disable=protected-access
--> 151           self._serialization_cache)
    152       self._functions[obj] = obj_functions
    153     if extra_functions:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in _list_functions_for_serialization(self, serialization_cache)
   2611     self.predict_function = None
   2612     functions = super(
-> 2613         Model, self)._list_functions_for_serialization(serialization_cache)
   2614     self.train_function = train_function
   2615     self.test_function = test_function

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in _list_functions_for_serialization(self, serialization_cache)
   3085   def _list_functions_for_serialization(self, serialization_cache):
   3086     return (self._trackable_saved_model_saver
-> 3087             .list_functions_for_serialization(serialization_cache))
   3088 
   3089   def __getstate__(self):

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/base_serialization.py in list_functions_for_serialization(self, serialization_cache)
     92       return {}
     93 
---> 94     fns = self.functions_to_serialize(serialization_cache)
     95 
     96     # The parent AutoTrackable class saves all user-defined tf.functions, and

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py in functions_to_serialize(self, serialization_cache)
     77   def functions_to_serialize(self, serialization_cache):
     78     return (self._get_serialized_attributes(
---> 79         serialization_cache).functions_to_serialize)
     80 
     81   def _get_serialized_attributes(self, serialization_cache):

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes(self, serialization_cache)
     93 
     94     object_dict, function_dict = self._get_serialized_attributes_internal(
---> 95         serialization_cache)
     96 
     97     serialized_attr.set_and_validate_objects(object_dict)

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/model_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
     55     objects, functions = (
     56         super(ModelSavedModelSaver, self)._get_serialized_attributes_internal(
---> 57             serialization_cache))
     58     functions['_default_save_signature'] = default_signature
     59     return objects, functions

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
    102     """Returns dictionary of serialized attributes."""
    103     objects = save_impl.wrap_layer_objects(self.obj, serialization_cache)
--> 104     functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
    105     # Attribute validator requires that the default save signature is added to
    106     # function dict, even if the value is None.

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in wrap_layer_functions(layer, serialization_cache)
    163   call_fn_with_losses = call_collection.add_function(
    164       _wrap_call_and_conditional_losses(layer),
--> 165       '{}_layer_call_and_return_conditional_losses'.format(layer.name))
    166   call_fn = call_collection.add_function(
    167       _extract_outputs_from_fn(layer, call_fn_with_losses),

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in add_function(self, call_fn, name)
    503       # Manually add traces for layers that have keyword arguments and have
    504       # a fully defined input signature.
--> 505       self.add_trace(*self._input_signature)
    506     return fn
    507 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in add_trace(self, *args, **kwargs)
    418             fn.get_concrete_function(*args, **kwargs)
    419 
--> 420         trace_with_training(True)
    421         trace_with_training(False)
    422       else:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in trace_with_training(value, fn)
    416           utils.set_training_arg(value, self._training_arg_index, args, kwargs)
    417           with K.deprecated_internal_learning_phase_scope(value):
--> 418             fn.get_concrete_function(*args, **kwargs)
    419 
    420         trace_with_training(True)

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in get_concrete_function(self, *args, **kwargs)
    548     if not self.call_collection.tracing:
    549       self.call_collection.add_trace(*args, **kwargs)
--> 550     return super(LayerCall, self).get_concrete_function(*args, **kwargs)
    551 
    552 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in get_concrete_function(self, *args, **kwargs)
   1297       ValueError: if this object has not yet been called on concrete values.
   1298     """
-> 1299     concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
   1300     concrete._garbage_collector.release()  # pylint: disable=protected-access
   1301     return concrete

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   1203       if self._stateful_fn is None:
   1204         initializers = []
-> 1205         self._initialize(args, kwargs, add_initializers_to=initializers)
   1206         self._initialize_uninitialized_variables(initializers)
   1207 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    724     self._concrete_stateful_fn = (
    725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
--> 726             *args, **kwds))
    727 
    728     def invalid_creator_scope(*unused_args, **unused_kwds):

/opt/conda/lib/python3.7/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 

/opt/conda/lib/python3.7/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 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3204             arg_names=arg_names,
   3205             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3206             capture_by_value=self._capture_by_value),
   3207         self._function_attributes,
   3208         function_spec=self.function_spec,

/opt/conda/lib/python3.7/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,

/opt/conda/lib/python3.7/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 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
    525       with autocast_variable.enable_auto_cast_variables(
    526           layer._compute_dtype_object):  # pylint: disable=protected-access
--> 527         ret = method(*args, **kwargs)
    528     _restore_layer_losses(original_losses)
    529     return ret

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    169     return control_flow_util.smart_cond(
    170         training, lambda: replace_training_and_call(True),
--> 171         lambda: replace_training_and_call(False))
    172 
    173   # Create arg spec for decorated function. If 'training' is not defined in the

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    113         pred, true_fn=true_fn, false_fn=false_fn, name=name)
    114   return smart_module.smart_cond(
--> 115       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    116 
    117 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     52   if pred_value is not None:
     53     if pred_value:
---> 54       return true_fn()
     55     else:
     56       return false_fn()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in <lambda>()
    168 
    169     return control_flow_util.smart_cond(
--> 170         training, lambda: replace_training_and_call(True),
    171         lambda: replace_training_and_call(False))
    172 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    165     def replace_training_and_call(training):
    166       set_training_arg(training, training_arg_index, args, kwargs)
--> 167       return wrapped_call(*args, **kwargs)
    168 
    169     return control_flow_util.smart_cond(

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(inputs, *args, **kwargs)
    568   def call_and_return_conditional_losses(inputs, *args, **kwargs):
    569     """Returns layer (call_output, conditional losses) tuple."""
--> 570     call_output = layer_call(inputs, *args, **kwargs)
    571     if version_utils.is_v1_layer_or_model(layer):
    572       conditional_losses = layer.get_losses_for(inputs)

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py in call(self, inputs, training, mask)
    423     """
    424     return self._run_internal_graph(
--> 425         inputs, training=training, mask=mask)
    426 
    427   def compute_output_shape(self, input_shape):

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py in _run_internal_graph(self, inputs, training, mask)
    558 
    559         args, kwargs = node.map_arguments(tensor_dict)
--> 560         outputs = node.layer(*args, **kwargs)
    561 
    562         # Update tensor_dict.

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in return_outputs_and_add_losses(*args, **kwargs)
     71     inputs = args[inputs_arg_index]
     72     args = args[inputs_arg_index + 1:]
---> 73     outputs, losses = fn(inputs, *args, **kwargs)
     74     layer.add_loss(losses, inputs=inputs)
     75 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    169     return control_flow_util.smart_cond(
    170         training, lambda: replace_training_and_call(True),
--> 171         lambda: replace_training_and_call(False))
    172 
    173   # Create arg spec for decorated function. If 'training' is not defined in the

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    113         pred, true_fn=true_fn, false_fn=false_fn, name=name)
    114   return smart_module.smart_cond(
--> 115       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    116 
    117 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     52   if pred_value is not None:
     53     if pred_value:
---> 54       return true_fn()
     55     else:
     56       return false_fn()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in <lambda>()
    168 
    169     return control_flow_util.smart_cond(
--> 170         training, lambda: replace_training_and_call(True),
    171         lambda: replace_training_and_call(False))
    172 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    165     def replace_training_and_call(training):
    166       set_training_arg(training, training_arg_index, args, kwargs)
--> 167       return wrapped_call(*args, **kwargs)
    168 
    169     return control_flow_util.smart_cond(

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in __call__(self, *args, **kwargs)
    542   def __call__(self, *args, **kwargs):
    543     if not self.call_collection.tracing:
--> 544       self.call_collection.add_trace(*args, **kwargs)
    545     return super(LayerCall, self).__call__(*args, **kwargs)
    546 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in add_trace(self, *args, **kwargs)
    418             fn.get_concrete_function(*args, **kwargs)
    419 
--> 420         trace_with_training(True)
    421         trace_with_training(False)
    422       else:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in trace_with_training(value, fn)
    416           utils.set_training_arg(value, self._training_arg_index, args, kwargs)
    417           with K.deprecated_internal_learning_phase_scope(value):
--> 418             fn.get_concrete_function(*args, **kwargs)
    419 
    420         trace_with_training(True)

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in get_concrete_function(self, *args, **kwargs)
    548     if not self.call_collection.tracing:
    549       self.call_collection.add_trace(*args, **kwargs)
--> 550     return super(LayerCall, self).get_concrete_function(*args, **kwargs)
    551 
    552 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in get_concrete_function(self, *args, **kwargs)
   1297       ValueError: if this object has not yet been called on concrete values.
   1298     """
-> 1299     concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
   1300     concrete._garbage_collector.release()  # pylint: disable=protected-access
   1301     return concrete

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   1215       # run the first trace but we should fail if variables are created.
   1216       concrete = self._stateful_fn._get_concrete_function_garbage_collected(  # pylint: disable=protected-access
-> 1217           *args, **kwargs)
   1218       if self._created_variables:
   1219         raise ValueError("Creating variables on a non-first call to a function"

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   3017       args, kwargs = None, None
   3018     with self._lock:
-> 3019       graph_function, _ = self._maybe_define_function(args, kwargs)
   3020       seen_names = set()
   3021       captured = object_identity.ObjectIdentitySet(

/opt/conda/lib/python3.7/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 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3204             arg_names=arg_names,
   3205             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3206             capture_by_value=self._capture_by_value),
   3207         self._function_attributes,
   3208         function_spec=self.function_spec,

/opt/conda/lib/python3.7/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,

/opt/conda/lib/python3.7/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 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
    525       with autocast_variable.enable_auto_cast_variables(
    526           layer._compute_dtype_object):  # pylint: disable=protected-access
--> 527         ret = method(*args, **kwargs)
    528     _restore_layer_losses(original_losses)
    529     return ret

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    169     return control_flow_util.smart_cond(
    170         training, lambda: replace_training_and_call(True),
--> 171         lambda: replace_training_and_call(False))
    172 
    173   # Create arg spec for decorated function. If 'training' is not defined in the

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    113         pred, true_fn=true_fn, false_fn=false_fn, name=name)
    114   return smart_module.smart_cond(
--> 115       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    116 
    117 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     52   if pred_value is not None:
     53     if pred_value:
---> 54       return true_fn()
     55     else:
     56       return false_fn()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in <lambda>()
    168 
    169     return control_flow_util.smart_cond(
--> 170         training, lambda: replace_training_and_call(True),
    171         lambda: replace_training_and_call(False))
    172 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    165     def replace_training_and_call(training):
    166       set_training_arg(training, training_arg_index, args, kwargs)
--> 167       return wrapped_call(*args, **kwargs)
    168 
    169     return control_flow_util.smart_cond(

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(inputs, *args, **kwargs)
    568   def call_and_return_conditional_losses(inputs, *args, **kwargs):
    569     """Returns layer (call_output, conditional losses) tuple."""
--> 570     call_output = layer_call(inputs, *args, **kwargs)
    571     if version_utils.is_v1_layer_or_model(layer):
    572       conditional_losses = layer.get_losses_for(inputs)

/opt/conda/lib/python3.7/site-packages/transformers/models/albert/modeling_tf_albert.py in call(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, labels, training, **kwargs)
   1144             labels=labels,
   1145             training=training,
-> 1146             kwargs_call=kwargs,
   1147         )
   1148         outputs = self.albert(

/opt/conda/lib/python3.7/site-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
    372                     output[tensor_name] = input
    373                 else:
--> 374                     output[parameter_names[i]] = input
    375             elif isinstance(input, allowed_types) or input is None:
    376                 output[parameter_names[i]] = input

IndexError: list index out of range
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

In version <= 4.0.0 use save method was without the error, but after `input_processing` function comment that the error happened.
Does any advice to fix the problem?



Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

I am having the same problem with BERT. My solution so far was to downgrade to `transformers==4.0.1` which seems to be the last version which does not use `input_processing` in `TFBertMainLayer`.

In my case, the values of the relevant variables are
```
input_ids = [<tf.Tensor 'input_ids:0' shape=(None, 384) dtype=int32>, <tf.Tensor 'input_ids_1:0' shape=(None, 384) dtype=int32>]
parameter_names = ['args']
```
The error arises because of the second item in `input_ids`. Like in the previous example I am using BERT as a part of a larger Keras model. Both the larger model and BERT have one input layer with name `input_ids`. I suspect that this is the reason why the list `input_ids` contains two elements. If I wrap `output[parameter_names[i]] = input` in a try-catch, it works as intended.