H5 to Pb Conversion with Fake Quantization Node Fails

Describe the bug

I have created a model in which I have used the annotate layer API to quantize some of the layers. The weights/model are saved in h5 format. I want the model and weights to be saved in pb format as I in my pipeline I have a requirement for that. When I try generating pb file I get the error.

System information

Tensorflow is installed from binary

TensorFlow version: 2.1.0

TensorFlow Model Optimization version: 0.3.0

Python version: 3.7.6

Describe the expected behavior

The model should generate the pb file with all of the fake quant layers intact.

Describe the current behavior

The pb file generation fails with error message -
ValueError: Attempted to save a function b'__inference_Conv2D1_layer_call_fn_585' which references a symbolic Tensor Tensor("model/quant_Conv2D1/LastValueQuant/FakeQuantWithMinMaxVarsPerChannel:0", shape=(5, 5, 3, 32), dtype=float32) that is not a simple constant. This is not supported.

Code to reproduce the issue

model.py

import numpy as np
from tensorflow.keras.layers import Input, Conv2D, Lambda
from tensorflow.keras.models import Model,load_model
import tensorflow as tf
import tensorflow_model_optimization as tfmot

class  espcn:
def  __init__(self, scale_factor=4, image_channels=3,loader=False):
self.__name__ = 'espcn'
self.scale_factor = scale_factor
self.channels = image_channels
self.loader = loader

# upsampling the resolution of image
def  sub_pixel(self, x):
return tf.compat.v1.depth_to_space(x, self.scale_factor,name="depth2space")
# building the espcn network

def  __call__(self):
if  self.loader is  True:
input_image = Input(shape=(240, 432, self.channels), name='x')
else:
input_image = Input(shape=(None, None, self.channels), name='x')\
x = tfmot.quantization.keras.quantize_annotate_layer(Conv2D(32, 5, kernel_initializer='glorot_uniform', padding='same', activation=tf.nn.relu,name="Conv2D1"))(input_image)
x = tfmot.quantization.keras.quantize_annotate_layer(Conv2D(32, 3, kernel_initializer='glorot_uniform', padding='same',activation=tf.nn.relu,name="Conv2D2"))(x)
x = tfmot.quantization.keras.quantize_annotate_layer(Conv2D(self.scale_factor**2*self.channels, 3, kernel_initializer='glorot_uniform', padding='same',activation=tf.nn.relu,name="Conv2D3"))(x)
if  self.scale_factor > 1:
#x = Lambda(self.sub_pixel)(x)
x = self.sub_pixel(x)
model = Model(inputs=input_image, outputs=x)
return model

train.py

if  __name__ == '__main__':
# get the espcn model
espcn = espcn_model.espcn(scale_factor=args.scale_factor)
model = espcn()
#q_aware stands for for quantization aware.
q_aware_model = tfmot.quantization.keras.quantize_apply(model)
model=q_aware_model
# print the network structure of model
model.sumAdd any other context about the problem here.mary()

model.compile(optimizer=Adam(args.lr), loss='mse')
#set keras callback function to save the model
checkpointer = ModelCheckpoint(os.path.join(save_dir,'model_weights_{epoch:03d}.h5'),
verbose=1, save_weights_only=False,save_freq='epoch', period=args.save_every)
# set keras callback function to dynamically adjust the learning rate
lr_scheduler = LearningRateScheduler(lr_schedule)

# start train 
model.fit(train_datagen(batch_size=args.batch_size), epochs=args.epoch, verbose=1, steps_per_epoch = 300,callbacks=[lr_scheduler,checkpointer])

convert.py

espcn = espcn_model.espcn(scale_factor=4,loader=True)
model = espcn()
q_aware_model = tfmot.quantization.keras.quantize_apply(model)
#Loading the 1000th epcoh weights
q_aware_model.load_weights('./models/espcn/run1/model_weights_995.hdf5')
model=q_aware_model
# print the network structure of model
model.summary()
#Printing the input and outputs of the model
print(model.outputs)
print(model.inputs)
#checkpoint_directory = "./saved_pb"
#checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
#checkpoint = tf.train.Checkpoint(model=model)
#checkpoint.save(file_prefix=checkpoint_prefix)
model.save('./saved_pb_new')

Additional context
The checkpoints file has no problem in saving but when I save the pb file it gives me error.

Start with convert.py and reproduce the failure at model.save('./saved_pb_new') using the listed TensorFlow 2.1.0 and TensorFlow Model Optimization 0.3.0 versions. Read model.py and train.py to trace how quantize_annotate_layer, quantize_apply, and the H5 checkpoint are used. Done means the model saves to PB format while retaining the fake quantization layers.