# Buggy code
def quantize_annotate_model(to_annotate):
    """Annotate a model to be quantized.

    This function does not actually quantize anything. It is merely to specify
    that the model needs to be quantized.

    This function is intended to be used in conjunction with the
    `quantize_annotate_layer`
    API. It's otherwise simpler to use `quantize_model`.

    Annotate a model while overriding the default behavior for one layer:

    ```python
    quantize_config = MyDenseQuantizeConfig()

    model = quantize_annotate_model(keras.Sequential([
        layers.Dense(10, activation='relu', input_shape=(100,)),
        quantize_annotate_layer(layers.Dense(2, activation='sigmoid'),
        quantize_config=quantize_config)
    ])))
    ```

    Note that this function removes the optimizer from the original model.

    Args:
      to_annotate: tf.keras model to annotate to be quantized.

    Returns:
      New tf.keras model with each layer in the model wrapped with
      `QuantizeAnnotate`.
    """
    if to_annotate is None:
        raise ValueError("`to_annotate` cannot be None")

    if not isinstance(to_annotate, keras.Model):
        raise ValueError(
            "`to_annotate` can only be a `tf.keras.Model` instance. Use "
            "the `quantize_annotate_layer` API to handle individual layers. "
            "You passed an instance of type: {input}.".format(
                input=to_annotate.__class__.__name__
            )
        )

    def _add_quant_wrapper(layer):
        # Already annotated layer. No need to wrap.
        if isinstance(layer, quantize_annotate_mod.QuantizeAnnotate):
            return layer

        return quantize_annotate_mod.QuantizeAnnotate(layer)

    return keras.models.clone_model(
        to_annotate, input_tensors=None, clone_function=_add_quant_wrapper
    )


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-ff8232533be0> in <module>
3 outputs = backbone(input_image)
4 not_working_model = keras.Model(input_image, outputs)
----> 5 q_aware_model = quantize_model(not_working_model)

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_model(to_quantize)
99         'the `quantize_annotate_layer` API to handle individual layers.')
100
--> 101   annotated_model = quantize_annotate_model(to_quantize)
102   return quantize_apply(annotated_model)
103

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_annotate_model(to_annotate)
148
149   return keras.models.clone_model(
--> 150       to_annotate, input_tensors=None, clone_function=_add_quant_wrapper)
151
152

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in clone_model(model, input_tensors, clone_function)
425   else:
426     return _clone_functional_model(
--> 427         model, input_tensors=input_tensors, layer_fn=clone_function)
428
429

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/models.py in _clone_functional_model(model, input_tensors, layer_fn)
198   input_tensors, output_tensors, created_layers = (
199       network.reconstruct_from_config(model_config,
--> 200                                       created_layers=created_layers))
201   metrics_names = model.metrics_names
202   model = Model(input_tensors, output_tensors, name=model.name)

~/anaconda3/envs/tensorflow2/lib/python3.7/site-packages/tensorflow/python/keras/engine/network.py in reconstruct_from_config(config, custom_objects, created_layers)
2044     layer = created_layers[layer_name]
2045     node_index = get_node_index(layer, node_index)
-> 2046     layer_output_tensors = layer._inbound_nodes[node_index].output_tensors
2047     output_tensors.append(nest.flatten(layer_output_tensors)[tensor_index])
2048

TypeError: list indices must be integers or slices, not NoneType

# Fixed code
def quantize_annotate_model(to_annotate):
    """Annotate a model to be quantized.

    This function does not actually quantize anything. It is merely to specify
    that the model needs to be quantized.

    This function is intended to be used in conjunction with the
    `quantize_annotate_layer`
    API. It's otherwise simpler to use `quantize_model`.

    Annotate a model while overriding the default behavior for one layer:

    ```python
    quantize_config = MyDenseQuantizeConfig()

    model = quantize_annotate_model(keras.Sequential([
        layers.Dense(10, activation='relu', input_shape=(100,)),
        quantize_annotate_layer(layers.Dense(2, activation='sigmoid'),
        quantize_config=quantize_config)
    ])))
    ```

    Note that this function removes the optimizer from the original model.

    Args:
      to_annotate: tf.keras model to annotate to be quantized.

    Returns:
      New tf.keras model with each layer in the model wrapped with
      `QuantizeAnnotate`.
    """
    if to_annotate is None:
        raise ValueError("`to_annotate` cannot be None")

    if not isinstance(to_annotate, keras.Model):
        raise ValueError(
            "`to_annotate` can only be a `tf.keras.Model` instance. Use "
            "the `quantize_annotate_layer` API to handle individual layers. "
            "You passed an instance of type: {input}.".format(
                input=to_annotate.__class__.__name__
            )
        )

    if (
        not isinstance(to_annotate, keras.Sequential)
        and not to_annotate._is_graph_network
    ):  # pylint: disable=protected-access
        raise ValueError(
            "`to_annotate` can only either be a tf.keras Sequential or "
            "Functional model."
        )

    def _add_quant_wrapper(layer):
        """Add annotation wrapper."""
        # Already annotated layer. No need to wrap.
        if isinstance(layer, quantize_annotate_mod.QuantizeAnnotate):
            return layer

        if isinstance(layer, tf.keras.Model):
            raise ValueError(
                "Quantizing a tf.keras Model inside another tf.keras Model is not supported."
            )

        return quantize_annotate_mod.QuantizeAnnotate(layer)

    return keras.models.clone_model(
        to_annotate, input_tensors=None, clone_function=_add_quant_wrapper
    )