How to save wrapped DistilBERT without using `save_pretrained`?
Author: hardianlawi
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.9.2
- Platform: Ubuntu 20
- Python version: 3.8
- PyTorch version (GPU?): 
- Tensorflow version (GPU?): 2.6
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- 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, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- encoder-decoder models (For example, BlenderBot, BART, Marian, Pegasus, T5, ByT5): @patrickvonplaten, @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj @patrickvonplaten 
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten
- Tokenizers: @LysandreJik
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

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

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

@Rocketknight1 

## Information

Model I am using (Bert, XLNet ...):

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)

## To reproduce

<!-- 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.-->

Simply run the codes below

```python
import tensorflow as tf
from transformers import (
    TFDistilBertModel,
    DistilBertTokenizerFast,
    DistilBertConfig,
)


def build_classifier_model():
    input_ids = tf.keras.layers.Input(shape=(None,), name="input_ids", dtype=tf.int32)
    attention_mask = tf.keras.layers.Input(
        shape=(None,), name="attention_mask", dtype=tf.int32
    )

    config = DistilBertConfig(
        dropout=0.2,
        attention_dropout=0.2,
        output_attentions=True,
        output_hidden_states=False,
        return_dict=False,
    )
    transformer = TFDistilBertModel.from_pretrained(
        "distilbert-base-uncased", config=config
    )
    transformer.trainable = False

    last_hidden_state = transformer(
        [input_ids, attention_mask],
    )[0]

    x = last_hidden_state[:, 0, :]
    x = tf.keras.layers.Dense(768, activation="relu")(x)
    x = tf.keras.layers.Dropout(0.2)(x)

    outputs = {
        label_name: tf.keras.layers.Dense(1, activation="sigmoid", name=label_name)(x)
        for label_name in ['A', 'B', 'C']
    }

    return tf.keras.Model([input_ids, attention_mask], outputs)

model = build_classifier_model()
model.save('./dump/savedmodel')
```

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

I expect this to generate artifacts containing the model in savedmodel format, but instead I got

```
~/miniforge3/envs/folder/lib/python3.8/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
    561         **kwargs,
    562     ):
--> 563         inputs = input_processing(
    564             func=self.call,
    565             config=self.config,

~/miniforge3/envs/folder/lib/python3.8/site-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
    376                     output[tensor_name] = input
    377                 else:
--> 378                     output[parameter_names[i]] = input
    379             elif isinstance(input, allowed_types) or input is None:
    380                 output[parameter_names[i]] = input

IndexError: list index out of range
```

Comments:
I saw others posted similar issues https://github.com/huggingface/transformers/issues/13610 and https://github.com/huggingface/transformers/issues/13742. However, since I am wrapping the model in `tf.keras.Model`, `save_pretrained` isn't a viable solution. Are there any workarounds?

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.

This issue should be resolved by recent PRs - if you're still encountering difficulties after installing the most recent release, please reopen it and let us know!

@hardianlawi  Were you able to solve this issue? I am still facing this issue, can you help? 

Hi @Rocketknight1  i am still facing this issue. I am not able to save finetuned `TFDistilBertModel` model in keras with `model.save()` . Since the model is wrapped in `tf.keras.Model` I can't use `save_pretrained`.
transformers version: 4.15
Platform: Ubuntu 20
Python version: 3.8
PyTorch version (GPU?):
Tensorflow version (GPU?): 2.6.2
Using GPU in script?: Yes
Using distributed or parallel set-up in script?: Yes

Error:
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_47167/2242234445.py in <module>
----> 1 model.save("save_path")

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/training.py in save(self, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
   2143     """
   2144     # pylint: enable=line-too-long
-> 2145     save.save_model(self, filepath, overwrite, include_optimizer, save_format,
   2146                     signatures, options, save_traces)
   2147 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/save.py in save_model(model, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
    147   else:
    148     with generic_utils.SharedObjectSavingScope():
--> 149       saved_model_save.save(model, filepath, overwrite, include_optimizer,
    150                             signatures, options, save_traces)
    151 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save.py in save(model, filepath, overwrite, include_optimizer, signatures, options, save_traces)
     88   with K.deprecated_internal_learning_phase_scope(0):
     89     with utils.keras_option_scope(save_traces):
---> 90       saved_nodes, node_paths = save_lib.save_and_return_nodes(
     91           model, filepath, signatures, options)
     92 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in save_and_return_nodes(obj, export_dir, signatures, options, experimental_skip_checkpoint)
   1226 
   1227   _, exported_graph, object_saver, asset_info, saved_nodes, node_paths = (
-> 1228       _build_meta_graph(obj, signatures, options, meta_graph_def))
   1229   saved_model.saved_model_schema_version = (
   1230       pywrap_libexport.SAVED_MODEL_SCHEMA_VERSION)

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph(obj, signatures, options, meta_graph_def)
   1397 
   1398   with save_context.save_context(options):
-> 1399     return _build_meta_graph_impl(obj, signatures, options, meta_graph_def)

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph_impl(obj, signatures, options, meta_graph_def)
   1333   checkpoint_graph_view = _AugmentedGraphView(obj)
   1334   if signatures is None:
-> 1335     signatures = signature_serialization.find_function_to_export(
   1336         checkpoint_graph_view)
   1337 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/signature_serialization.py in find_function_to_export(saveable_view)
     97   # If the user did not specify signatures, check the root object for a function
     98   # that can be made into a signature.
---> 99   functions = saveable_view.list_functions(saveable_view.root)
    100   signature = functions.get(DEFAULT_SIGNATURE_ATTR, None)
    101   if signature is not None:

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in list_functions(self, obj)
    161     obj_functions = self._functions.get(obj, None)
    162     if obj_functions is None:
--> 163       obj_functions = obj._list_functions_for_serialization(  # pylint: disable=protected-access
    164           self._serialization_cache)
    165       self._functions[obj] = obj_functions

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/training.py in _list_functions_for_serialization(self, serialization_cache)
   2810     self.predict_function = None
   2811     self.train_tf_function = None
-> 2812     functions = super(
   2813         Model, self)._list_functions_for_serialization(serialization_cache)
   2814     self.train_function = train_function

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/base_layer.py in _list_functions_for_serialization(self, serialization_cache)
   3083 
   3084   def _list_functions_for_serialization(self, serialization_cache):
-> 3085     return (self._trackable_saved_model_saver
   3086             .list_functions_for_serialization(serialization_cache))
   3087 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/base_serialization.py in list_functions_for_serialization(self, serialization_cache)
     91       return {}
     92 
---> 93     fns = self.functions_to_serialize(serialization_cache)
     94 
     95     # The parent AutoTrackable class saves all user-defined tf.functions, and

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in functions_to_serialize(self, serialization_cache)
     71 
     72   def functions_to_serialize(self, serialization_cache):
---> 73     return (self._get_serialized_attributes(
     74         serialization_cache).functions_to_serialize)
     75 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes(self, serialization_cache)
     87       return serialized_attr
     88 
---> 89     object_dict, function_dict = self._get_serialized_attributes_internal(
     90         serialization_cache)
     91 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/model_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
     54     # the ones serialized by Layer.
     55     objects, functions = (
---> 56         super(ModelSavedModelSaver, self)._get_serialized_attributes_internal(
     57             serialization_cache))
     58     functions['_default_save_signature'] = default_signature

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
     97     """Returns dictionary of serialized attributes."""
     98     objects = save_impl.wrap_layer_objects(self.obj, serialization_cache)
---> 99     functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
    100     # Attribute validator requires that the default save signature is added to
    101     # function dict, even if the value is None.

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrap_layer_functions(layer, serialization_cache)
    195       for fn in fns.values():
    196         if fn is not None and not isinstance(fn, LayerCall):
--> 197           fn.get_concrete_function()
    198 
    199   # Restore overwritten functions and losses

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/contextlib.py in __exit__(self, type, value, traceback)
    118         if type is None:
    119             try:
--> 120                 next(self.gen)
    121             except StopIteration:
    122                 return False

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in tracing_scope()
    357       if training is not None:
    358         with K.deprecated_internal_learning_phase_scope(training):
--> 359           fn.get_concrete_function(*args, **kwargs)
    360       else:
    361         fn.get_concrete_function(*args, **kwargs)

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in get_concrete_function(self, *args, **kwargs)
   1231   def get_concrete_function(self, *args, **kwargs):
   1232     # Implements GenericFunction.get_concrete_function.
-> 1233     concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
   1234     concrete._garbage_collector.release()  # pylint: disable=protected-access
   1235     return concrete

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   1211       if self._stateful_fn is None:
   1212         initializers = []
-> 1213         self._initialize(args, kwargs, add_initializers_to=initializers)
   1214         self._initialize_uninitialized_variables(initializers)
   1215 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    757     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    758     self._concrete_stateful_fn = (
--> 759         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    760             *args, **kwds))
    761 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   3064       args, kwargs = None, None
   3065     with self._lock:
-> 3066       graph_function, _ = self._maybe_define_function(args, kwargs)
   3067     return graph_function
   3068 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3461 
   3462           self._function_cache.missed.add(call_context_key)
-> 3463           graph_function = self._create_graph_function(args, kwargs)
   3464           self._function_cache.primary[cache_key] = graph_function
   3465 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3296     arg_names = base_arg_names + missing_arg_names
   3297     graph_function = ConcreteFunction(
-> 3298         func_graph_module.func_graph_from_py_func(
   3299             self._name,
   3300             self._python_function,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
   1005         _, original_func = tf_decorator.unwrap(python_func)
   1006 
-> 1007       func_outputs = python_func(*func_args, **func_kwargs)
   1008 
   1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    666         # the function a weak reference to itself to avoid a reference cycle.
    667         with OptionalXlaContext(compile_with_xla):
--> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    669         return out
    670 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
    570       with autocast_variable.enable_auto_cast_variables(
    571           layer._compute_dtype_object):  # pylint: disable=protected-access
--> 572         ret = method(*args, **kwargs)
    573     _restore_layer_losses(original_losses)
    574     return ret

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    162       return wrapped_call(*args, **kwargs)
    163 
--> 164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
    166         lambda: replace_training_and_call(False))

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    103     return tf.cond(
    104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
--> 105   return tf.__internal__.smart_cond.smart_cond(
    106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    107 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     56       return true_fn()
     57     else:
---> 58       return false_fn()
     59   else:
     60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
    164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
--> 166         lambda: replace_training_and_call(False))
    167 
    168   # Create arg spec for decorated function. If 'training' is not defined in the

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    160     def replace_training_and_call(training):
    161       set_training_arg(training, training_arg_index, args, kwargs)
--> 162       return wrapped_call(*args, **kwargs)
    163 
    164     return control_flow_util.smart_cond(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call(inputs, *args, **kwargs)
    649     return layer.keras_api.__call__  # pylint: disable=protected-access
    650   def call(inputs, *args, **kwargs):
--> 651     return call_and_return_conditional_losses(inputs, *args, **kwargs)[0]
    652   return _create_call_fn_decorator(layer, call)
    653 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in __call__(self, *args, **kwargs)
    607   def __call__(self, *args, **kwargs):
    608     self._maybe_trace(args, kwargs)
--> 609     return self.wrapped_call(*args, **kwargs)
    610 
    611   def get_concrete_function(self, *args, **kwargs):

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    883 
    884       with OptionalXlaContext(self._jit_compile):
--> 885         result = self._call(*args, **kwds)
    886 
    887       new_tracing_count = self.experimental_get_tracing_count()

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    931       # This is the first call of __call__, so we have to initialize.
    932       initializers = []
--> 933       self._initialize(args, kwds, add_initializers_to=initializers)
    934     finally:
    935       # At this point we know that the initialization is complete (or less

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    757     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    758     self._concrete_stateful_fn = (
--> 759         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    760             *args, **kwds))
    761 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   3064       args, kwargs = None, None
   3065     with self._lock:
-> 3066       graph_function, _ = self._maybe_define_function(args, kwargs)
   3067     return graph_function
   3068 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3461 
   3462           self._function_cache.missed.add(call_context_key)
-> 3463           graph_function = self._create_graph_function(args, kwargs)
   3464           self._function_cache.primary[cache_key] = graph_function
   3465 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3296     arg_names = base_arg_names + missing_arg_names
   3297     graph_function = ConcreteFunction(
-> 3298         func_graph_module.func_graph_from_py_func(
   3299             self._name,
   3300             self._python_function,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
   1005         _, original_func = tf_decorator.unwrap(python_func)
   1006 
-> 1007       func_outputs = python_func(*func_args, **func_kwargs)
   1008 
   1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    666         # the function a weak reference to itself to avoid a reference cycle.
    667         with OptionalXlaContext(compile_with_xla):
--> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    669         return out
    670 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
    570       with autocast_variable.enable_auto_cast_variables(
    571           layer._compute_dtype_object):  # pylint: disable=protected-access
--> 572         ret = method(*args, **kwargs)
    573     _restore_layer_losses(original_losses)
    574     return ret

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    162       return wrapped_call(*args, **kwargs)
    163 
--> 164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
    166         lambda: replace_training_and_call(False))

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    103     return tf.cond(
    104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
--> 105   return tf.__internal__.smart_cond.smart_cond(
    106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    107 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     56       return true_fn()
     57     else:
---> 58       return false_fn()
     59   else:
     60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
    164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
--> 166         lambda: replace_training_and_call(False))
    167 
    168   # Create arg spec for decorated function. If 'training' is not defined in the

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    160     def replace_training_and_call(training):
    161       set_training_arg(training, training_arg_index, args, kwargs)
--> 162       return wrapped_call(*args, **kwargs)
    163 
    164     return control_flow_util.smart_cond(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(*args, **kwargs)
    631   def call_and_return_conditional_losses(*args, **kwargs):
    632     """Returns layer (call_output, conditional losses) tuple."""
--> 633     call_output = layer_call(*args, **kwargs)
    634     if version_utils.is_v1_layer_or_model(layer):
    635       conditional_losses = layer.get_losses_for(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/functional.py in call(self, inputs, training, mask)
    412         a list of tensors if there are more than one outputs.
    413     """
--> 414     return self._run_internal_graph(
    415         inputs, training=training, mask=mask)
    416 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/functional.py in _run_internal_graph(self, inputs, training, mask)
    548 
    549         args, kwargs = node.map_arguments(tensor_dict)
--> 550         outputs = node.layer(*args, **kwargs)
    551 
    552         # Update tensor_dict.

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in return_outputs_and_add_losses(*args, **kwargs)
     66       args = args[1:]
     67 
---> 68     outputs, losses = fn(*args, **kwargs)
     69     layer.add_loss(losses, inputs=True)
     70 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    162       return wrapped_call(*args, **kwargs)
    163 
--> 164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
    166         lambda: replace_training_and_call(False))

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    103     return tf.cond(
    104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
--> 105   return tf.__internal__.smart_cond.smart_cond(
    106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    107 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     56       return true_fn()
     57     else:
---> 58       return false_fn()
     59   else:
     60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
    164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
--> 166         lambda: replace_training_and_call(False))
    167 
    168   # Create arg spec for decorated function. If 'training' is not defined in the

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    160     def replace_training_and_call(training):
    161       set_training_arg(training, training_arg_index, args, kwargs)
--> 162       return wrapped_call(*args, **kwargs)
    163 
    164     return control_flow_util.smart_cond(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in __call__(self, *args, **kwargs)
    607   def __call__(self, *args, **kwargs):
    608     self._maybe_trace(args, kwargs)
--> 609     return self.wrapped_call(*args, **kwargs)
    610 
    611   def get_concrete_function(self, *args, **kwargs):

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    883 
    884       with OptionalXlaContext(self._jit_compile):
--> 885         result = self._call(*args, **kwds)
    886 
    887       new_tracing_count = self.experimental_get_tracing_count()

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    922       # In this case we have not created variables on the first call. So we can
    923       # run the first trace but we should fail if variables are created.
--> 924       results = self._stateful_fn(*args, **kwds)
    925       if self._created_variables and not ALLOW_DYNAMIC_VARIABLE_CREATION:
    926         raise ValueError("Creating variables on a non-first call to a function"

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in __call__(self, *args, **kwargs)
   3036     with self._lock:
   3037       (graph_function,
-> 3038        filtered_flat_args) = self._maybe_define_function(args, kwargs)
   3039     return graph_function._call_flat(
   3040         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3461 
   3462           self._function_cache.missed.add(call_context_key)
-> 3463           graph_function = self._create_graph_function(args, kwargs)
   3464           self._function_cache.primary[cache_key] = graph_function
   3465 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3296     arg_names = base_arg_names + missing_arg_names
   3297     graph_function = ConcreteFunction(
-> 3298         func_graph_module.func_graph_from_py_func(
   3299             self._name,
   3300             self._python_function,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
   1005         _, original_func = tf_decorator.unwrap(python_func)
   1006 
-> 1007       func_outputs = python_func(*func_args, **func_kwargs)
   1008 
   1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    666         # the function a weak reference to itself to avoid a reference cycle.
    667         with OptionalXlaContext(compile_with_xla):
--> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    669         return out
    670 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
    570       with autocast_variable.enable_auto_cast_variables(
    571           layer._compute_dtype_object):  # pylint: disable=protected-access
--> 572         ret = method(*args, **kwargs)
    573     _restore_layer_losses(original_losses)
    574     return ret

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
    162       return wrapped_call(*args, **kwargs)
    163 
--> 164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
    166         lambda: replace_training_and_call(False))

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
    103     return tf.cond(
    104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
--> 105   return tf.__internal__.smart_cond.smart_cond(
    106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
    107 

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
     56       return true_fn()
     57     else:
---> 58       return false_fn()
     59   else:
     60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
    164     return control_flow_util.smart_cond(
    165         training, lambda: replace_training_and_call(True),
--> 166         lambda: replace_training_and_call(False))
    167 
    168   # Create arg spec for decorated function. If 'training' is not defined in the

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
    160     def replace_training_and_call(training):
    161       set_training_arg(training, training_arg_index, args, kwargs)
--> 162       return wrapped_call(*args, **kwargs)
    163 
    164     return control_flow_util.smart_cond(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(*args, **kwargs)
    631   def call_and_return_conditional_losses(*args, **kwargs):
    632     """Returns layer (call_output, conditional losses) tuple."""
--> 633     call_output = layer_call(*args, **kwargs)
    634     if version_utils.is_v1_layer_or_model(layer):
    635       conditional_losses = layer.get_losses_for(

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
    560         **kwargs,
    561     ):
--> 562         inputs = input_processing(
    563             func=self.call,
    564             config=self.config,

~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
    418                     output[tensor_name] = input
    419                 else:
--> 420                     output[parameter_names[i]] = input
    421             elif isinstance(input, allowed_types) or input is None:
    422                 output[parameter_names[i]] = input

IndexError: list index out of range
```


Attaching code to replicate

```
import os

import tensorflow as tf
from tensorflow import keras
from keras import backend as K
from transformers import TFDistilBertModel, DistilBertConfig
from focal_loss import SparseCategoricalFocalLoss

MAX_LENGTH = 256
LAYER_DROPOUT = 0.2
LEARNING_RATE = 5e-5
RANDOM_STATE = 42
NUM_CLASSES=3


# Compatible with tensorflow backend

def focal_loss(gamma=2., alpha=.25):
    def focal_loss_fixed(y_true, y_pred):
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
        return -K.mean(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1+K.epsilon())) - K.mean((1 - alpha) * K.pow(pt_0, gamma) * K.log(1. - pt_0 + K.epsilon()))
    return focal_loss_fixed

def build_model(transformer, max_length=MAX_LENGTH):

    # Define weight initializer with a random seed to ensure reproducibility
    weight_initializer = tf.keras.initializers.GlorotNormal(seed=RANDOM_STATE) 
    
    # Define input layers
    input_ids_layer = tf.keras.layers.Input(shape=(max_length,), 
                                            name='input_ids', 
                                            dtype='int32')
    input_attention_layer = tf.keras.layers.Input(shape=(max_length,), 
                                                  name='attention_mask', 
                                                  dtype='int32')
#     input_attention_layer = tf.keras.layers.Input(shape=(max_length,), 
#                                                   name='attention_mask', 
#                                                   dtype='int32')
    
    # Extract [CLS] embedding
    # It is a tf.Tensor of shape (batch_size, sequence_length, hidden_size=768).
    last_hidden_state = transformer([input_ids_layer, input_attention_layer])[0]
    cls_token = last_hidden_state[:, 0, :]
    
    ##                                                 ##
    ## Define additional dropout and dense layers here ##
    ##                                                 ##
    
    # Define a FCN layer
    output = tf.keras.layers.Dense(NUM_CLASSES, 
                                   activation='softmax',
                                   kernel_initializer=weight_initializer,  
                                   kernel_constraint=None,
                                   bias_initializer='zeros'
                                   )(cls_token)
    
    # Define the model
#     {"input_ids": input_ids}
    model = tf.keras.Model([input_ids_layer, input_attention_layer], output)
    
    # Compile the model
    model.compile(tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE), 
                  loss=SparseCategoricalFocalLoss(gamma=2),
                  metrics=['accuracy'])
    
    return model

    


def get_distil_bert_model(trainable=False, config=None):
    if not config:
        DISTILBERT_DROPOUT = 0.2
        DISTILBERT_ATT_DROPOUT = 0.2

        # Configure DistilBERT's initialization
        config = DistilBertConfig(dropout=DISTILBERT_DROPOUT, 
                                  attention_dropout=DISTILBERT_ATT_DROPOUT, 
                                  output_hidden_states=False)

    distilBert = TFDistilBertModel.from_pretrained('distilbert-base-uncased', config=config)
     
    if trainable is False:
        for layer in distilBert.layers:
            layer.trainable = False
            
    return distilBert
    
def get_compiled_model():
    distilBert=get_distil_bert_model()
    classification_model=build_model(distilBert)
    return classification_model

model=get_compiled_model()
model.save("model_save_path")
```

> @hardianlawi Were you able to solve this issue? I am still facing this issue, can you help?
> 
> Hi @Rocketknight1 i am still facing this issue. I am not able to save finetuned `TFDistilBertModel` model in keras with `model.save()` . Since the model is wrapped in `tf.keras.Model` I can't use `save_pretrained`. transformers version: 4.15 Platform: Ubuntu 20 Python version: 3.8 PyTorch version (GPU?): Tensorflow version (GPU?): 2.6.2 Using GPU in script?: Yes Using distributed or parallel set-up in script?: Yes
> 
> Error:
> 
> ```
> ---------------------------------------------------------------------------
> IndexError                                Traceback (most recent call last)
> /tmp/ipykernel_47167/2242234445.py in <module>
> ----> 1 model.save("save_path")
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/training.py in save(self, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
>    2143     """
>    2144     # pylint: enable=line-too-long
> -> 2145     save.save_model(self, filepath, overwrite, include_optimizer, save_format,
>    2146                     signatures, options, save_traces)
>    2147 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/save.py in save_model(model, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
>     147   else:
>     148     with generic_utils.SharedObjectSavingScope():
> --> 149       saved_model_save.save(model, filepath, overwrite, include_optimizer,
>     150                             signatures, options, save_traces)
>     151 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save.py in save(model, filepath, overwrite, include_optimizer, signatures, options, save_traces)
>      88   with K.deprecated_internal_learning_phase_scope(0):
>      89     with utils.keras_option_scope(save_traces):
> ---> 90       saved_nodes, node_paths = save_lib.save_and_return_nodes(
>      91           model, filepath, signatures, options)
>      92 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in save_and_return_nodes(obj, export_dir, signatures, options, experimental_skip_checkpoint)
>    1226 
>    1227   _, exported_graph, object_saver, asset_info, saved_nodes, node_paths = (
> -> 1228       _build_meta_graph(obj, signatures, options, meta_graph_def))
>    1229   saved_model.saved_model_schema_version = (
>    1230       pywrap_libexport.SAVED_MODEL_SCHEMA_VERSION)
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph(obj, signatures, options, meta_graph_def)
>    1397 
>    1398   with save_context.save_context(options):
> -> 1399     return _build_meta_graph_impl(obj, signatures, options, meta_graph_def)
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph_impl(obj, signatures, options, meta_graph_def)
>    1333   checkpoint_graph_view = _AugmentedGraphView(obj)
>    1334   if signatures is None:
> -> 1335     signatures = signature_serialization.find_function_to_export(
>    1336         checkpoint_graph_view)
>    1337 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/signature_serialization.py in find_function_to_export(saveable_view)
>      97   # If the user did not specify signatures, check the root object for a function
>      98   # that can be made into a signature.
> ---> 99   functions = saveable_view.list_functions(saveable_view.root)
>     100   signature = functions.get(DEFAULT_SIGNATURE_ATTR, None)
>     101   if signature is not None:
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py in list_functions(self, obj)
>     161     obj_functions = self._functions.get(obj, None)
>     162     if obj_functions is None:
> --> 163       obj_functions = obj._list_functions_for_serialization(  # pylint: disable=protected-access
>     164           self._serialization_cache)
>     165       self._functions[obj] = obj_functions
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/training.py in _list_functions_for_serialization(self, serialization_cache)
>    2810     self.predict_function = None
>    2811     self.train_tf_function = None
> -> 2812     functions = super(
>    2813         Model, self)._list_functions_for_serialization(serialization_cache)
>    2814     self.train_function = train_function
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/base_layer.py in _list_functions_for_serialization(self, serialization_cache)
>    3083 
>    3084   def _list_functions_for_serialization(self, serialization_cache):
> -> 3085     return (self._trackable_saved_model_saver
>    3086             .list_functions_for_serialization(serialization_cache))
>    3087 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/base_serialization.py in list_functions_for_serialization(self, serialization_cache)
>      91       return {}
>      92 
> ---> 93     fns = self.functions_to_serialize(serialization_cache)
>      94 
>      95     # The parent AutoTrackable class saves all user-defined tf.functions, and
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in functions_to_serialize(self, serialization_cache)
>      71 
>      72   def functions_to_serialize(self, serialization_cache):
> ---> 73     return (self._get_serialized_attributes(
>      74         serialization_cache).functions_to_serialize)
>      75 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes(self, serialization_cache)
>      87       return serialized_attr
>      88 
> ---> 89     object_dict, function_dict = self._get_serialized_attributes_internal(
>      90         serialization_cache)
>      91 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/model_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
>      54     # the ones serialized by Layer.
>      55     objects, functions = (
> ---> 56         super(ModelSavedModelSaver, self)._get_serialized_attributes_internal(
>      57             serialization_cache))
>      58     functions['_default_save_signature'] = default_signature
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/layer_serialization.py in _get_serialized_attributes_internal(self, serialization_cache)
>      97     """Returns dictionary of serialized attributes."""
>      98     objects = save_impl.wrap_layer_objects(self.obj, serialization_cache)
> ---> 99     functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
>     100     # Attribute validator requires that the default save signature is added to
>     101     # function dict, even if the value is None.
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrap_layer_functions(layer, serialization_cache)
>     195       for fn in fns.values():
>     196         if fn is not None and not isinstance(fn, LayerCall):
> --> 197           fn.get_concrete_function()
>     198 
>     199   # Restore overwritten functions and losses
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/contextlib.py in __exit__(self, type, value, traceback)
>     118         if type is None:
>     119             try:
> --> 120                 next(self.gen)
>     121             except StopIteration:
>     122                 return False
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in tracing_scope()
>     357       if training is not None:
>     358         with K.deprecated_internal_learning_phase_scope(training):
> --> 359           fn.get_concrete_function(*args, **kwargs)
>     360       else:
>     361         fn.get_concrete_function(*args, **kwargs)
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in get_concrete_function(self, *args, **kwargs)
>    1231   def get_concrete_function(self, *args, **kwargs):
>    1232     # Implements GenericFunction.get_concrete_function.
> -> 1233     concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
>    1234     concrete._garbage_collector.release()  # pylint: disable=protected-access
>    1235     return concrete
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
>    1211       if self._stateful_fn is None:
>    1212         initializers = []
> -> 1213         self._initialize(args, kwargs, add_initializers_to=initializers)
>    1214         self._initialize_uninitialized_variables(initializers)
>    1215 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
>     757     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
>     758     self._concrete_stateful_fn = (
> --> 759         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
>     760             *args, **kwds))
>     761 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
>    3064       args, kwargs = None, None
>    3065     with self._lock:
> -> 3066       graph_function, _ = self._maybe_define_function(args, kwargs)
>    3067     return graph_function
>    3068 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
>    3461 
>    3462           self._function_cache.missed.add(call_context_key)
> -> 3463           graph_function = self._create_graph_function(args, kwargs)
>    3464           self._function_cache.primary[cache_key] = graph_function
>    3465 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
>    3296     arg_names = base_arg_names + missing_arg_names
>    3297     graph_function = ConcreteFunction(
> -> 3298         func_graph_module.func_graph_from_py_func(
>    3299             self._name,
>    3300             self._python_function,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
>    1005         _, original_func = tf_decorator.unwrap(python_func)
>    1006 
> -> 1007       func_outputs = python_func(*func_args, **func_kwargs)
>    1008 
>    1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
>     666         # the function a weak reference to itself to avoid a reference cycle.
>     667         with OptionalXlaContext(compile_with_xla):
> --> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
>     669         return out
>     670 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
>     570       with autocast_variable.enable_auto_cast_variables(
>     571           layer._compute_dtype_object):  # pylint: disable=protected-access
> --> 572         ret = method(*args, **kwargs)
>     573     _restore_layer_losses(original_losses)
>     574     return ret
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
>     162       return wrapped_call(*args, **kwargs)
>     163 
> --> 164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
>     166         lambda: replace_training_and_call(False))
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
>     103     return tf.cond(
>     104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
> --> 105   return tf.__internal__.smart_cond.smart_cond(
>     106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
>     107 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
>      56       return true_fn()
>      57     else:
> ---> 58       return false_fn()
>      59   else:
>      60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
>     164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
> --> 166         lambda: replace_training_and_call(False))
>     167 
>     168   # Create arg spec for decorated function. If 'training' is not defined in the
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
>     160     def replace_training_and_call(training):
>     161       set_training_arg(training, training_arg_index, args, kwargs)
> --> 162       return wrapped_call(*args, **kwargs)
>     163 
>     164     return control_flow_util.smart_cond(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call(inputs, *args, **kwargs)
>     649     return layer.keras_api.__call__  # pylint: disable=protected-access
>     650   def call(inputs, *args, **kwargs):
> --> 651     return call_and_return_conditional_losses(inputs, *args, **kwargs)[0]
>     652   return _create_call_fn_decorator(layer, call)
>     653 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in __call__(self, *args, **kwargs)
>     607   def __call__(self, *args, **kwargs):
>     608     self._maybe_trace(args, kwargs)
> --> 609     return self.wrapped_call(*args, **kwargs)
>     610 
>     611   def get_concrete_function(self, *args, **kwargs):
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
>     883 
>     884       with OptionalXlaContext(self._jit_compile):
> --> 885         result = self._call(*args, **kwds)
>     886 
>     887       new_tracing_count = self.experimental_get_tracing_count()
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
>     931       # This is the first call of __call__, so we have to initialize.
>     932       initializers = []
> --> 933       self._initialize(args, kwds, add_initializers_to=initializers)
>     934     finally:
>     935       # At this point we know that the initialization is complete (or less
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
>     757     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
>     758     self._concrete_stateful_fn = (
> --> 759         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
>     760             *args, **kwds))
>     761 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
>    3064       args, kwargs = None, None
>    3065     with self._lock:
> -> 3066       graph_function, _ = self._maybe_define_function(args, kwargs)
>    3067     return graph_function
>    3068 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
>    3461 
>    3462           self._function_cache.missed.add(call_context_key)
> -> 3463           graph_function = self._create_graph_function(args, kwargs)
>    3464           self._function_cache.primary[cache_key] = graph_function
>    3465 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
>    3296     arg_names = base_arg_names + missing_arg_names
>    3297     graph_function = ConcreteFunction(
> -> 3298         func_graph_module.func_graph_from_py_func(
>    3299             self._name,
>    3300             self._python_function,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
>    1005         _, original_func = tf_decorator.unwrap(python_func)
>    1006 
> -> 1007       func_outputs = python_func(*func_args, **func_kwargs)
>    1008 
>    1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
>     666         # the function a weak reference to itself to avoid a reference cycle.
>     667         with OptionalXlaContext(compile_with_xla):
> --> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
>     669         return out
>     670 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
>     570       with autocast_variable.enable_auto_cast_variables(
>     571           layer._compute_dtype_object):  # pylint: disable=protected-access
> --> 572         ret = method(*args, **kwargs)
>     573     _restore_layer_losses(original_losses)
>     574     return ret
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
>     162       return wrapped_call(*args, **kwargs)
>     163 
> --> 164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
>     166         lambda: replace_training_and_call(False))
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
>     103     return tf.cond(
>     104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
> --> 105   return tf.__internal__.smart_cond.smart_cond(
>     106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
>     107 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
>      56       return true_fn()
>      57     else:
> ---> 58       return false_fn()
>      59   else:
>      60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
>     164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
> --> 166         lambda: replace_training_and_call(False))
>     167 
>     168   # Create arg spec for decorated function. If 'training' is not defined in the
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
>     160     def replace_training_and_call(training):
>     161       set_training_arg(training, training_arg_index, args, kwargs)
> --> 162       return wrapped_call(*args, **kwargs)
>     163 
>     164     return control_flow_util.smart_cond(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(*args, **kwargs)
>     631   def call_and_return_conditional_losses(*args, **kwargs):
>     632     """Returns layer (call_output, conditional losses) tuple."""
> --> 633     call_output = layer_call(*args, **kwargs)
>     634     if version_utils.is_v1_layer_or_model(layer):
>     635       conditional_losses = layer.get_losses_for(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/functional.py in call(self, inputs, training, mask)
>     412         a list of tensors if there are more than one outputs.
>     413     """
> --> 414     return self._run_internal_graph(
>     415         inputs, training=training, mask=mask)
>     416 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/functional.py in _run_internal_graph(self, inputs, training, mask)
>     548 
>     549         args, kwargs = node.map_arguments(tensor_dict)
> --> 550         outputs = node.layer(*args, **kwargs)
>     551 
>     552         # Update tensor_dict.
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
>    1035         with autocast_variable.enable_auto_cast_variables(
>    1036             self._compute_dtype_object):
> -> 1037           outputs = call_fn(inputs, *args, **kwargs)
>    1038 
>    1039         if self._activity_regularizer:
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in return_outputs_and_add_losses(*args, **kwargs)
>      66       args = args[1:]
>      67 
> ---> 68     outputs, losses = fn(*args, **kwargs)
>      69     layer.add_loss(losses, inputs=True)
>      70 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
>     162       return wrapped_call(*args, **kwargs)
>     163 
> --> 164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
>     166         lambda: replace_training_and_call(False))
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
>     103     return tf.cond(
>     104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
> --> 105   return tf.__internal__.smart_cond.smart_cond(
>     106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
>     107 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
>      56       return true_fn()
>      57     else:
> ---> 58       return false_fn()
>      59   else:
>      60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
>     164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
> --> 166         lambda: replace_training_and_call(False))
>     167 
>     168   # Create arg spec for decorated function. If 'training' is not defined in the
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
>     160     def replace_training_and_call(training):
>     161       set_training_arg(training, training_arg_index, args, kwargs)
> --> 162       return wrapped_call(*args, **kwargs)
>     163 
>     164     return control_flow_util.smart_cond(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in __call__(self, *args, **kwargs)
>     607   def __call__(self, *args, **kwargs):
>     608     self._maybe_trace(args, kwargs)
> --> 609     return self.wrapped_call(*args, **kwargs)
>     610 
>     611   def get_concrete_function(self, *args, **kwargs):
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
>     883 
>     884       with OptionalXlaContext(self._jit_compile):
> --> 885         result = self._call(*args, **kwds)
>     886 
>     887       new_tracing_count = self.experimental_get_tracing_count()
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
>     922       # In this case we have not created variables on the first call. So we can
>     923       # run the first trace but we should fail if variables are created.
> --> 924       results = self._stateful_fn(*args, **kwds)
>     925       if self._created_variables and not ALLOW_DYNAMIC_VARIABLE_CREATION:
>     926         raise ValueError("Creating variables on a non-first call to a function"
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in __call__(self, *args, **kwargs)
>    3036     with self._lock:
>    3037       (graph_function,
> -> 3038        filtered_flat_args) = self._maybe_define_function(args, kwargs)
>    3039     return graph_function._call_flat(
>    3040         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
>    3461 
>    3462           self._function_cache.missed.add(call_context_key)
> -> 3463           graph_function = self._create_graph_function(args, kwargs)
>    3464           self._function_cache.primary[cache_key] = graph_function
>    3465 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
>    3296     arg_names = base_arg_names + missing_arg_names
>    3297     graph_function = ConcreteFunction(
> -> 3298         func_graph_module.func_graph_from_py_func(
>    3299             self._name,
>    3300             self._python_function,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/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, acd_record_initial_resource_uses)
>    1005         _, original_func = tf_decorator.unwrap(python_func)
>    1006 
> -> 1007       func_outputs = python_func(*func_args, **func_kwargs)
>    1008 
>    1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
>     666         # the function a weak reference to itself to avoid a reference cycle.
>     667         with OptionalXlaContext(compile_with_xla):
> --> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
>     669         return out
>     670 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in wrapper(*args, **kwargs)
>     570       with autocast_variable.enable_auto_cast_variables(
>     571           layer._compute_dtype_object):  # pylint: disable=protected-access
> --> 572         ret = method(*args, **kwargs)
>     573     _restore_layer_losses(original_losses)
>     574     return ret
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in wrap_with_training_arg(*args, **kwargs)
>     162       return wrapped_call(*args, **kwargs)
>     163 
> --> 164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
>     166         lambda: replace_training_and_call(False))
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/utils/control_flow_util.py in smart_cond(pred, true_fn, false_fn, name)
>     103     return tf.cond(
>     104         pred, true_fn=true_fn, false_fn=false_fn, name=name)
> --> 105   return tf.__internal__.smart_cond.smart_cond(
>     106       pred, true_fn=true_fn, false_fn=false_fn, name=name)
>     107 
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/tensorflow/python/framework/smart_cond.py in smart_cond(pred, true_fn, false_fn, name)
>      56       return true_fn()
>      57     else:
> ---> 58       return false_fn()
>      59   else:
>      60     return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in <lambda>()
>     164     return control_flow_util.smart_cond(
>     165         training, lambda: replace_training_and_call(True),
> --> 166         lambda: replace_training_and_call(False))
>     167 
>     168   # Create arg spec for decorated function. If 'training' is not defined in the
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/utils.py in replace_training_and_call(training)
>     160     def replace_training_and_call(training):
>     161       set_training_arg(training, training_arg_index, args, kwargs)
> --> 162       return wrapped_call(*args, **kwargs)
>     163 
>     164     return control_flow_util.smart_cond(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/keras/saving/saved_model/save_impl.py in call_and_return_conditional_losses(*args, **kwargs)
>     631   def call_and_return_conditional_losses(*args, **kwargs):
>     632     """Returns layer (call_output, conditional losses) tuple."""
> --> 633     call_output = layer_call(*args, **kwargs)
>     634     if version_utils.is_v1_layer_or_model(layer):
>     635       conditional_losses = layer.get_losses_for(
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
>     560         **kwargs,
>     561     ):
> --> 562         inputs = input_processing(
>     563             func=self.call,
>     564             config=self.config,
> 
> ~/SageMaker/custom-miniconda/miniconda/envs/custom_python_38/lib/python3.8/site-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
>     418                     output[tensor_name] = input
>     419                 else:
> --> 420                     output[parameter_names[i]] = input
>     421             elif isinstance(input, allowed_types) or input is None:
>     422                 output[parameter_names[i]] = input
> 
> IndexError: list index out of range
> ```
> 
> Attaching code to replicate
> 
> ```
> import os
> 
> import tensorflow as tf
> from tensorflow import keras
> from keras import backend as K
> from transformers import TFDistilBertModel, DistilBertConfig
> from focal_loss import SparseCategoricalFocalLoss
> 
> MAX_LENGTH = 256
> LAYER_DROPOUT = 0.2
> LEARNING_RATE = 5e-5
> RANDOM_STATE = 42
> NUM_CLASSES=3
> 
> 
> # Compatible with tensorflow backend
> 
> def focal_loss(gamma=2., alpha=.25):
>     def focal_loss_fixed(y_true, y_pred):
>         pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
>         pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
>         return -K.mean(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1+K.epsilon())) - K.mean((1 - alpha) * K.pow(pt_0, gamma) * K.log(1. - pt_0 + K.epsilon()))
>     return focal_loss_fixed
> 
> def build_model(transformer, max_length=MAX_LENGTH):
> 
>     # Define weight initializer with a random seed to ensure reproducibility
>     weight_initializer = tf.keras.initializers.GlorotNormal(seed=RANDOM_STATE) 
>     
>     # Define input layers
>     input_ids_layer = tf.keras.layers.Input(shape=(max_length,), 
>                                             name='input_ids', 
>                                             dtype='int32')
>     input_attention_layer = tf.keras.layers.Input(shape=(max_length,), 
>                                                   name='attention_mask', 
>                                                   dtype='int32')
> #     input_attention_layer = tf.keras.layers.Input(shape=(max_length,), 
> #                                                   name='attention_mask', 
> #                                                   dtype='int32')
>     
>     # Extract [CLS] embedding
>     # It is a tf.Tensor of shape (batch_size, sequence_length, hidden_size=768).
>     last_hidden_state = transformer([input_ids_layer, input_attention_layer])[0]
>     cls_token = last_hidden_state[:, 0, :]
>     
>     ##                                                 ##
>     ## Define additional dropout and dense layers here ##
>     ##                                                 ##
>     
>     # Define a FCN layer
>     output = tf.keras.layers.Dense(NUM_CLASSES, 
>                                    activation='softmax',
>                                    kernel_initializer=weight_initializer,  
>                                    kernel_constraint=None,
>                                    bias_initializer='zeros'
>                                    )(cls_token)
>     
>     # Define the model
> #     {"input_ids": input_ids}
>     model = tf.keras.Model([input_ids_layer, input_attention_layer], output)
>     
>     # Compile the model
>     model.compile(tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE), 
>                   loss=SparseCategoricalFocalLoss(gamma=2),
>                   metrics=['accuracy'])
>     
>     return model
> 
>     
> 
> 
> def get_distil_bert_model(trainable=False, config=None):
>     if not config:
>         DISTILBERT_DROPOUT = 0.2
>         DISTILBERT_ATT_DROPOUT = 0.2
> 
>         # Configure DistilBERT's initialization
>         config = DistilBertConfig(dropout=DISTILBERT_DROPOUT, 
>                                   attention_dropout=DISTILBERT_ATT_DROPOUT, 
>                                   output_hidden_states=False)
> 
>     distilBert = TFDistilBertModel.from_pretrained('distilbert-base-uncased', config=config)
>      
>     if trainable is False:
>         for layer in distilBert.layers:
>             layer.trainable = False
>             
>     return distilBert
>     
> def get_compiled_model():
>     distilBert=get_distil_bert_model()
>     classification_model=build_model(distilBert)
>     return classification_model
> 
> model=get_compiled_model()
> model.save("model_save_path")
> ```


I have the same problem, How to solve it?



@kapilkd13 @Zjq9409 I completely switched to Pytorch and Pytorch Lightning since they made my life easier :')