# Buggy code
def quantize_model(to_quantize):
    """Quantize a whole tf.keras model with the default quantization implementation.

    To be more precise, `quantize_model` creates a model that emulates
    quantization during training and stores information that downstream
    tools will use to produce actually quantized models.

    Quantize a model:

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

    Note that this function removes the optimizer from the original model.
    Additionally, training the model returned by `quantize_model` will not affect
    the weights of the original model.

    Args:
      to_quantize: tf.keras model to be quantized. It can have pre-trained
        weights.

    Returns:
      Returns a new tf.keras model prepared for quantization.
    """
    if to_quantize is None:
        raise ValueError("`to_quantize` cannot be None")

    if not isinstance(to_quantize, keras.Model):
        raise ValueError(
            "`to_quantize` 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_quantize.__class__.__name__
            )
        )

    annotated_model = quantize_annotate_model(to_quantize)
    return quantize_apply(annotated_model)


# 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_model(to_quantize):
    """Quantize a whole tf.keras model with the default quantization implementation.

    To be more precise, `quantize_model` creates a model that emulates
    quantization during training and stores information that downstream
    tools will use to produce actually quantized models.

    Quantize a model:

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

    Note that this function removes the optimizer from the original model.
    Additionally, training the model returned by `quantize_model` will not affect
    the weights of the original model.

    Args:
      to_quantize: tf.keras model to be quantized. It can have pre-trained
        weights.

    Returns:
      Returns a new tf.keras model prepared for quantization.
    """
    if to_quantize is None:
        raise ValueError("`to_quantize` cannot be None")

    if not isinstance(to_quantize, keras.Model):
        raise ValueError(
            "`to_quantize` 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_quantize.__class__.__name__
            )
        )

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

    annotated_model = quantize_annotate_model(to_quantize)
    return quantize_apply(annotated_model)