# Buggy code
def __init__(self, layer, quantize_config=None, **kwargs):
    """Create a quantize annotate wrapper over a keras layer.

    Args:
      layer: The keras layer to be quantized.
      quantize_config: `QuantizeConfig` to quantize layer.
      **kwargs: Additional keyword arguments to be passed to the keras layer.
    """
    super(QuantizeAnnotate, self).__init__(layer, **kwargs)

    self.quantize_config = quantize_config

    self._track_trackable(layer, name="layer")
    # Enables end-user to annotate the first layer in Sequential models, while
    # passing the input shape to the original layer.
    #
    # tf.keras.Sequential(
    #   quantize_annotate_layer(tf.keras.layers.Dense(2, input_shape=(3,)))
    # )
    #
    # as opposed to
    #
    # tf.keras.Sequential(
    #   quantize_annotate_layer(tf.keras.layers.Dense(2), input_shape=(3,))
    # )
    #
    # Without this code, the QuantizeAnnotate wrapper doesn't have an input
    # shape and being the first layer, this causes the model to not be
    # built. Being not built is confusing since the end-user has passed an
    # input shape.
    if not hasattr(self, "_batch_input_shape") and hasattr(layer, "_batch_input_shape"):
        self._batch_input_shape = self.layer._batch_input_shape  # pylint: disable=protected-access


# 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 __init__(self, layer, quantize_config=None, **kwargs):
    """Create a quantize annotate wrapper over a keras layer.

    Args:
      layer: The keras layer to be quantized.
      quantize_config: `QuantizeConfig` to quantize layer.
      **kwargs: Additional keyword arguments to be passed to the keras layer.
    """
    super(QuantizeAnnotate, self).__init__(layer, **kwargs)

    if layer is None:
        raise ValueError("`layer` cannot be None.")

    # Check against keras.Model since it is an instance of keras.layers.Layer.
    if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
        layer, tf.keras.Model
    ):
        raise ValueError(
            "`layer` can only be a `tf.keras.layers.Layer` instance. "
            "You passed an instance of type: {input}.".format(
                input=layer.__class__.__name__
            )
        )

    self.quantize_config = quantize_config

    self._track_trackable(layer, name="layer")
    # Enables end-user to annotate the first layer in Sequential models, while
    # passing the input shape to the original layer.
    #
    # tf.keras.Sequential(
    #   quantize_annotate_layer(tf.keras.layers.Dense(2, input_shape=(3,)))
    # )
    #
    # as opposed to
    #
    # tf.keras.Sequential(
    #   quantize_annotate_layer(tf.keras.layers.Dense(2), input_shape=(3,))
    # )
    #
    # Without this code, the QuantizeAnnotate wrapper doesn't have an input
    # shape and being the first layer, this causes the model to not be
    # built. Being not built is confusing since the end-user has passed an
    # input shape.
    if not hasattr(self, "_batch_input_shape") and hasattr(layer, "_batch_input_shape"):
        self._batch_input_shape = self.layer._batch_input_shape  # pylint: disable=protected-access