Assertion error with custom quantize config with per_axis=True

Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.

Describe the bug
I am trying to perform quantization-aware training on my model, with a custom quantize configuration with per-channel quantization. However, I get an assertion error.

System information

TensorFlow version (installed from source or binary): 2.4.1

TensorFlow Model Optimization version (installed from source or binary): 0.5.0

Python version: 3.6

Describe the expected behavior

model is correctly quantized

Describe the current behavior

tensorflow_model_optimization/python/core/quantization/keras/quant_ops.py:332 _FakeQuantWithMinMaxVars *
assert len(min_var.get_shape()) == 1
Assertion error

Code to reproduce the issue
Provide a reproducible code that is the bare minimum necessary to generate the
problem.

class MyQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):

  def get_weights_and_quantizers(self, layer):
    return [(layer.kernel, tfmot.quantization.keras.quantizers.LastValueQuantizer(num_bits=8, per_axis=True, symmetric=True, narrow_range=True))]

  def get_activations_and_quantizers(self, layer):
    return [(layer.activation, tfmot.quantization.keras.quantizers.MovingAverageQuantizer(num_bits=8, per_axis=True, symmetric=True, narrow_range=True))]

  def set_quantize_weights(self, layer, quantize_weights):
    # Add this line for each item returned in `get_weights_and_quantizers`
    # , in the same order
    layer.kernel = quantize_weights[0]

  def set_quantize_activations(self, layer, quantize_activations):
    # Add this line for each item returned in `get_activations_and_quantizers`
    # , in the same order.
    layer.activation = quantize_activations[0]

  def get_output_quantizers(self, layer):
    # Does not quantize output, since we return an empty list.
    return []

  def get_config(self):
    return {}

def custom_quantization(layer):
    if not isinstance(layer, tf.keras.layers.Dense):
        return tfmot.quantization.keras.quantize_annotate_layer(layer, MyQuantizeConfig())
    return layer


model = tf.keras.models.load_model(load_ckpt_path)
annotated_model = tf.keras.models.clone_model(model, clone_function=custom_quantization)
with tfmot.quantization.keras.quantize_scope({'MyQuantizeConfig': MyQuantizeConfig}):
    model = tfmot.quantization.keras.quantize_apply(annotated_model)

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
I cannot share my model, however it is a convolutional network with 2D convolutions, relu6 activations, a traditional inception block, batch normalizations and a final dense layer for classification.