Activity Regularizer not working with quantization aware training (QAT) TF 2.6

Describe the Issue
Activity Regularizer not working with quantization aware training (QAT). TypeError: An op outside of the function building code is being passed a "Graph" tensor.

System information

TensorFlow version (installed from source or binary): TF 2.6

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

Python version: 3.6.9

Describe the expected behavior
This happens only on quantized models, non-quantized model have numpy value for activity regularization loss. The model losses i.e. regularization losses should be accumulated without empty tensors/placeholders. And should be able to operate on them outside the graph when executing in eager mode.

Describe the current behavior
Currently the activity regularization loss tensor is of type <class 'tensorflow.python.framework.ops.Tensor'> with no numpy values, the bias regularization tensor is also tf.tensor but that does carry a numpy value. But I am not sure if that is throwing the error, where I am unable to operate on the the tf op.

Code to reproduce the issue

import numpy as np
import tensorflow as tf
from tensorflow_model_optimization.python.core.quantization.keras import quantize
from tensorflow.python import keras

l = tf.keras.layers

tf.config.run_functions_eagerly(True)

def functional_model():
    """Builds an MNIST functional model."""
    inp = tf.keras.Input(shape=image_input_shape())
    x = l.Conv2D(filters=32, kernel_size=5, padding='same', activation='relu',
                 kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(inp)
    x = tf.keras.layers.ActivityRegularization(l2=0.0001)(x)
    x = l.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same')(x)
    # TODO(pulkitb): Add BatchNorm when transformations are ready.
    # x = l.BatchNormalization()(x)
    x = l.Conv2D(filters=64, kernel_size=5, padding='same', activation='relu',
                 kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(x)
    x = tf.keras.layers.ActivityRegularization(l2=0.0001)(x)    
    x = l.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same')(x)
    x = l.Flatten()(x)
    x = l.Dense(1024, activation='relu')(x)
    x = l.Dropout(0.4)(x)
    out = l.Dense(10, activation='softmax')(x)

    return tf.keras.Model(inp, [out])


def image_input_shape(img_rows=28, img_cols=28):
    if tf.keras.backend.image_data_format() == 'channels_first':
        return 1, img_rows, img_cols
    else:
        return img_rows, img_cols, 1


def preprocessed_data(img_rows=28,
                      img_cols=28,
                      num_classes=10):
    """Get data for mnist training and evaluation."""
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

    if tf.keras.backend.image_data_format() == 'channels_first':
        x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
        x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
    else:
        x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
        x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255

    # convert class vectors to binary class matrices
    y_train = tf.keras.utils.to_categorical(y_train, num_classes)
    y_test = tf.keras.utils.to_categorical(y_test, num_classes)

    return x_train, y_train, x_test, y_test


model = functional_model()
model.summary()
x_train, y_train, x_test, y_test = preprocessed_data()

model.compile(
    loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
model.fit(x_train, y_train, batch_size=500)
_, model_accuracy = model.evaluate(x_test, y_test, verbose=0)

print("Quantizing model")

quantized_model = quantize.quantize_model(model)
print(quantized_model.losses)
quantized_model.compile(
    loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
print(quantized_model.losses)

quantized_model.fit(x_train, y_train, batch_size=500)
_, quantized_model_accuracy = quantized_model.evaluate(
    x_test, y_test, verbose=0)

Error Output

  120/120 [==============================] - 3s 19ms/step - loss: 2.1668 - accuracy: 0.3804
Quantizing model
[<tf.Tensor: shape=(), dtype=float32, numpy=0.00022464147>, <tf.Tensor 'activity_regularization/ActivityRegularizer/truediv:0' shape=() dtype=float32>, <tf.Tensor: shape=(), dtype=float32, numpy=0.0042893193>, <tf.Tensor 'activity_regularization_1/ActivityRegularizer/truediv:0' shape=() dtype=float32>]
[<tf.Tensor: shape=(), dtype=float32, numpy=0.00022464147>, <tf.Tensor 'activity_regularization/ActivityRegularizer/truediv:0' shape=() dtype=float32>, <tf.Tensor: shape=(), dtype=float32, numpy=0.0042893193>, <tf.Tensor 'activity_regularization_1/ActivityRegularizer/truediv:0' shape=() dtype=float32>]
Traceback (most recent call last):
  File "/home/user/git/archive-code/tempTrain.py", line 80, in <module>
    quantized_model.fit(x_train, y_train, batch_size=500)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/training.py", line 1184, in fit
    tmp_logs = self.train_function(iterator)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/training.py", line 853, in train_function
    return step_function(self, iterator)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/training.py", line 842, in step_function
    outputs = model.distribute_strategy.run(run_step, args=(data,))
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/distribute/distribute_lib.py", line 1286, in run
    return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/distribute/distribute_lib.py", line 2849, in call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/distribute/distribute_lib.py", line 3632, in _call_for_each_replica
    return fn(*args, **kwargs)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 597, in wrapper
    return func(*args, **kwargs)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/training.py", line 835, in run_step
    outputs = model.train_step(data)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/training.py", line 789, in train_step
    y, y_pred, sample_weight, regularization_losses=self.losses)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/keras/engine/compile_utils.py", line 231, in __call__
    reg_loss = tf.add_n(regularization_losses)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py", line 4000, in add_n
    return gen_math_ops.add_n(inputs, name=name)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py", line 392, in add_n
    inputs, name=name, ctx=_ctx)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py", line 425, in add_n_eager_fallback
    ctx=ctx, name=name)
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 75, in quick_execute
    raise e
  File "/home/user/TF2_6Venv/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: activity_regularization/ActivityRegularizer/truediv:0