title: Cannot access regularization loss values when performing quantization aware training (QAT) in tensorflow 2.3
body: <p>Cannot access regularization loss values when performing quantization aware training (QAT) in tensorflow 2.3.</p>
  <blockquote>
  <p>AttributeError: 'Tensor' object has no attribute 'numpy'</p>
  </blockquote>
  <p>I have created a test mnist model and added quantization nodes in the model using the TFMOT QAT library. While trying to access the values of regularization loss that gets collected in the <code>model.losses</code> as mentioned in the tensorflow examples, I get the above error. The values can be accessed fine when I am running the same without quantization (QAT). What could be causing this? I am using a functional model and performing custom training operation.</p>
  <p><strong>Code to reproduce the issue</strong></p>
  <pre><code>import numpy as np
  import tensorflow as tf
  from tensorflow_model_optimization.python.core.quantization.keras import quantize
  from tensorflow.python import keras
  from tensorflow.python.keras.utils import losses_utils
  l = keras.layers
  
  tf.config.run_functions_eagerly(True)
  physical_devices = tf.config.list_physical_devices('GPU')
  for device in physical_devices:
      tf.config.experimental.set_memory_growth(device, True)
  
  def functional_model():
    &quot;&quot;&quot;Builds an MNIST functional model.&quot;&quot;&quot;
    inp = keras.Input(image_input_shape())
    x = l.Conv2D(32, 5, padding='same', activation='relu', activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(inp)
    x = l.MaxPooling2D((2, 2), (2, 2), padding='same')(x)
    # TODO(pulkitb): Add BatchNorm when transformations are ready.
    # x = l.BatchNormalization()(x)
    x = l.Conv2D(64, 5, padding='same', activation='relu', activity_regularizer=tf.keras.regularizers.l2(l=0.0001), kernel_regularizer=tf.keras.regularizers.l2(l=0.0001))(x)
    x = l.MaxPooling2D((2, 2), (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 keras.models.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):
    &quot;&quot;&quot;Get data for mnist training and evaluation.&quot;&quot;&quot;
    (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
  
  @tf.function
  def trainStep(model, optimizer, images, labels):
      with tf.GradientTape() as tape:
          nnRunOutput = model(images, training=True)
          regTotalLoss = 0.0
          if model.losses:
              regLoss = losses_utils.cast_losses_to_common_dtype(model.losses)
              regTotalLoss = tf.math.add_n(regLoss)
          loss = tf.cast(tf.keras.losses.CategoricalCrossentropy(from_logits=True)(labels, nnRunOutput), dtype=tf.float32)
          totalLossTensor = tf.math.add_n([loss, regTotalLoss])
      gradients = tape.gradient(totalLossTensor, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
      print(loss.numpy())
      print(totalLossTensor.numpy())
      print(tf.math.add_n(regTotalLoss).numpy())
      trainLossItems = [loss.numpy(), tf.math.add_n(regTotalLoss).numpy()]
      print(trainLossItems)
  
  
  with tf.device('/device:GPU:' + &quot;0&quot;):
      model = functional_model() #sequential_model()
      model.summary()
      x_train, y_train, x_test, y_test = preprocessed_data()
      print(&quot;Quantizing model&quot;)
  
      quantized_model = quantize.quantize_model(model)
      print(quantized_model.losses)
  
  
      optimizer = tf.keras.optimizers.SGD()
      trainImgBatches = np.array_split(x_train, len(x_train)/500)
      trainLabelBatches = np.array_split(y_train, len(y_train)/500)
      for idx in range(0,len(trainImgBatches)):
          trainStep(quantized_model, optimizer, trainImgBatches[idx], trainLabelBatches[idx])
  </code></pre>
  <p><strong>Error output</strong></p>
  <pre><code>Quantizing model
  [&lt;tf.Tensor 'conv2d/ActivityRegularizer_2/truediv:0' shape=() dtype=float32&gt;, &lt;tf.Tensor: shape=(), dtype=float32, numpy=0.00019313334&gt;, &lt;tf.Tensor 'conv2d_1/ActivityRegularizer_2/truediv:0' shape=() dtype=float32&gt;, &lt;tf.Tensor: shape=(), dtype=float32, numpy=0.0042751725&gt;]
  2021-09-28 12:39:32.366940: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.7
  2021-09-28 12:39:32.947998: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
  2.3025882
  Traceback (most recent call last):
    File &quot;/home/user/git/archive-code/tempTrain.py&quot;, line 219, in &lt;module&gt;
      trainStep(quantized_model, optimizer, trainImgBatches[idx], trainLabelBatches[idx])
    File &quot;/home/user/TF2MainEnv/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py&quot;, line 761, in __call__
      return self._python_function(*args, **kwds)
    File &quot;/home/user/git/archive-code/tempTrain.py&quot;, line 186, in trainStep
      print(totalLossTensor.numpy())
  AttributeError: 'Tensor' object has no attribute 'numpy'
  
  Process finished with exit code 1
  </code></pre>