# Buggy code
def __init__(self, layer, quantize_config, **kwargs):
    """Create a quantize emulate wrapper for 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.
    """

    if quantize_config is None:
        raise ValueError(
            "quantize_config cannot be None. It is needed to quantize a layer."
        )

    if "name" not in kwargs:
        kwargs["name"] = self._make_layer_name(layer)

    super(QuantizeWrapper, self).__init__(layer, **kwargs)
    self.quantize_config = quantize_config

    self._track_trackable(layer, name="layer")


# 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, **kwargs):
    """Create a quantize emulate wrapper for 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.
    """
    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__
            )
        )

    if quantize_config is None:
        raise ValueError(
            "quantize_config cannot be None. It is needed to quantize a layer."
        )

    if "name" not in kwargs:
        kwargs["name"] = self._make_layer_name(layer)

    super(QuantizeWrapper, self).__init__(layer, **kwargs)
    self.quantize_config = quantize_config

    self._track_trackable(layer, name="layer")