Basically current QAT API doesn’t supports subclass model with just a quantize_model API.
I’d like to recommend you change the model to be functional model if the model is simple enough. (it may requires redo normal training to make sure the model is exactly same.)
Functional model Example) models/classification_model.py at f8f4845cc85ef674d6285337b55e43638039ff91 · tensorflow/models · GitHub 5
super().init(inputs=inputs, outputs=x) is a patten that you can create a subclass functional model which is internally same as an example here: A API Funcional  |  TensorFlow Core 2
But it also not supports to quantize recursively as Bhack@ mentioned.
So you may have to put a flag to init to determine you quantize the self.base model manually (by calling quantize_apply) on init method. (also you have to call quantize_apply for the entire model to quantize some layers outside of self.base.) As similar to the Bhack@ mentioned link.