RuntimeError: Layer backbone.stage0.0.fuse_layersb2.1.0.0.1:<class 'tensorflow.python.keras.layers.normalization_v2.BatchNormalization'> is not supported. You can quantize this layer by passing a `tfmot.quantization.keras.QuantizeConfig` instance to the `quantize_annotate_layer` API.

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
The batchnorm is not fused with the preceding conv2d and depth conv2d layers in the fuse_layer block of the code attached below and hence it is shown as an unsupported layer for quantization aware training

System information

TensorFlow version (installed from source or binary): tensorflow-gpu 2.4.0 install using pip

TensorFlow Model Optimization version (installed from source or binary):tensorflow-model-optimization 0.6.0 installed using pip

Python version: 3.7

Describe the expected behavior
the conv2d and depthwise conv2d must be fused with batchnorm which is followed after that so that the layers are ready for QAT

Describe the current behavior
the layers do not seem to be fused causing the above runtime error

Code to reproduce the issue

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras import layers
from tensorflow.keras.layers import Lambda


def spatial_weighting(input,channels,ratio=16,prefix='',i=''):

    x = tf.keras.layers.GlobalAveragePooling2D()(input)
    x = Lambda(lambda x:tf.expand_dims(x, axis=1))(x)
    x = Lambda(lambda x:tf.expand_dims(x, axis=1))(x)
    # conv1
    x = tf.keras.layers.Conv2D(int(channels / ratio),kernel_size=1,strides=1,use_bias=True,name=prefix+'.spatial_weighting.'+i+'.conv1.conv')(x)
    x = tf.keras.layers.Activation("relu")(x)

    # conv2
    x = tf.keras.layers.Conv2D(channels,kernel_size=1,strides=1,use_bias=True,name=prefix+'.spatial_weighting.'+i+'.conv2.conv')(x)
    x = tf.keras.layers.Activation("sigmoid")(x)
    x = tf.keras.layers.Multiply()([input,x])
    return x
    


def CrossResolutionWeighting(input,channels,
                 ratio=16,prefix=''):
        def get_avgpool_s_k_sz(x):
            outputsz = np.array(x[-1].shape[-3:-1])
            stridesz_list = []
            kernelsz_lsit = []
            for index in range(len(x)):
                inputsz_i = np.array(x[index].shape[-3:-1])
                stridesz_i = np.floor(inputsz_i / outputsz).astype(np.int32)
                kernelsz_i = inputsz_i - (outputsz - 1) * stridesz_i
                stridesz_list.append(stridesz_i)
                kernelsz_lsit.append(kernelsz_i)
            return stridesz_list,kernelsz_lsit,outputsz
        total_channel = sum(channels)
        mini_size = input[-1].shape[-3:-1].as_list()
        stridesz_list, kernelsz_list, outputsz = get_avgpool_s_k_sz(input)
        x = [tf.keras.layers.AveragePooling2D(pool_size=kernelsz.tolist(),strides=stridesz.tolist())(s) for s,kernelsz,stridesz in
               zip(input[:-1],kernelsz_list,stridesz_list) ] + [input[-1]]
        x = Lambda( lambda x: tf.concat(x,axis=3))(x)
        #conv1
        x = tf.keras.layers.Conv2D(int(total_channel / ratio),kernel_size=1,strides=1,use_bias=False,name=prefix+'.cross_resolution_weighting.conv1.conv')(x)
        x = tf.keras.layers.BatchNormalization(name=prefix+'.cross_resolution_weighting.conv1.bn')(x)
        x = tf.keras.layers.Activation("relu")(x)

        #conv2
        x = tf.keras.layers.Conv2D(total_channel,kernel_size=1,strides=1,use_bias=False,name=prefix+'.cross_resolution_weighting.conv2.conv')(x)
        x = tf.keras.layers.BatchNormalization(name=prefix+'.cross_resolution_weighting.conv2.bn')(x)
        x = tf.keras.layers.Activation("sigmoid")(x)

        x = Lambda( lambda x: tf.split(x,num_or_size_splits=channels,axis=3))(x)
        x1 = [
            Lambda( lambda x: tf.compat.v1.image.resize(x[0], size=x[1].shape[-3:-1], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR))([a,s])
            for s, a in zip(input, x)
        ]

        x = [tf.keras.layers.Multiply()([a,s]) for s,a in zip(input,x1)]
        return x

def depth_wise_unit(input,strides,channel,prefix='',i=''):
    x = tf.keras.layers.ZeroPadding2D(padding=1)(input)
    if prefix =='':
        x = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=strides,use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        return x
    elif i=='':
        x = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=strides,use_bias=False,name=prefix+'depthwise_conv.conv')(x)
        x = tf.keras.layers.BatchNormalization(name=prefix+'depthwise_conv.bn')(x)
        return x
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=strides,use_bias=False,name=prefix+'.depthwise_convs.'+str(i)+'.conv')(x)
    x = tf.keras.layers.BatchNormalization(name=prefix+'.depthwise_convs.'+str(i)+'.bn')(x)
    return x


def channel_shuffle(x, groups):
    """Channel Shuffle operation.

    This function enables cross-group information flow for multiple groups
    convolution layers.

    Args:
        x (Tensor): The input tensor.
        groups (int): The number of groups to divide the input tensor
            in the channel dimension.

    Returns:
        Tensor: The output tensor after channel shuffle operation.
    """

    batch_size, height, width, num_channels = x.shape
    
    assert (num_channels % groups == 0), ('num_channels should be '
                                          'divisible by groups')
    channels_per_group = num_channels // groups

    if batch_size is None:
        x =  tf.keras.layers.Reshape((height, width ,groups, channels_per_group))(x)
    else:
        x =  tf.keras.layers.Reshape((batch_size, height, width, groups, channels_per_group))(x)
    x = Lambda( lambda x: tf.transpose(x,perm=[0,1,2,4,3]))(x)
    if batch_size is None:
        x = tf.keras.layers.Reshape((height, width, -1))(x)
    else:
        x = tf.keras.layers.Reshape((batch_size, height, width, -1))(x)
    return x


def ConditionalChannelWeighting(input,in_channels,
                 strides,
                 reduce_ratio,prefix=''):
        assert strides in [1, 2]

        branch_channels = [channel // 2 for channel in in_channels]

        x = []
        for s in input:
            x.append(Lambda( lambda x: tf.split(x,num_or_size_splits=2,axis=3))(s))

        x1 = [s[0] for s in x]
        x2 = [s[1] for s in x]

        # CrossResolutionWeighting
        x2 = CrossResolutionWeighting(x2,branch_channels,ratio=reduce_ratio,prefix=prefix)

        #calling functions 
        x2 = [depth_wise_unit(s,strides,channel,prefix=prefix,i=str(i)) for i, (s, channel) in enumerate(zip(x2, branch_channels))]
        x2 = [spatial_weighting(s,channel,ratio=4,prefix=prefix,i=str(i)) for i, (s, channel) in enumerate(zip(x2, branch_channels))]

        x = []
        for s1,s2 in zip(x1,x2):
            x.append(Lambda( lambda x: tf.concat([x[0],x[1]],axis=3))([s1, s2]))
        x = [channel_shuffle(s, 2) for s in x]
        
        return x

def Stem(input, stem_channels,
                 out_channels,
                 expand_ratio):
    
    mid_channels = int(round(stem_channels * expand_ratio))
    branch_channels = stem_channels // 2
    if stem_channels == out_channels:
        inc_channels = out_channels - branch_channels
    else:
        inc_channels = out_channels - stem_channels
    # conv1
    x = tf.keras.layers.ZeroPadding2D(padding=1)(input)
    x = tf.keras.layers.Conv2D(stem_channels,kernel_size=3,strides=2,use_bias=False,name='backbone.stem.conv1.conv')(x)
    x = tf.keras.layers.BatchNormalization(name='backbone.stem.conv1.bn')(x)
    x = tf.keras.layers.Activation("relu")(x)
    x1, x2 = Lambda( lambda x: tf.split(x,num_or_size_splits=2,axis=3))(x)

    # expand conv
    x2 = tf.keras.layers.Conv2D(mid_channels,kernel_size=1,strides=1,use_bias=False,name='backbone.stem.expand_conv.conv')(x2) #,padding=0
    x2 = tf.keras.layers.BatchNormalization(name='backbone.stem.expand_conv.bn')(x2)
    x2 = tf.keras.layers.Activation("relu")(x2)
    
    # depthwise conv
    x2 = depth_wise_unit(x2,strides=2,channel=mid_channels,prefix='backbone.stem.')

    #linear conv
    if stem_channels == out_channels:
        x2 = tf.keras.layers.Conv2D(branch_channels,kernel_size=1,strides=1,use_bias=False,name='backbone.stem.linear_conv.conv')(x2) 
    else :
        x2 = tf.keras.layers.Conv2D(stem_channels,kernel_size=1,strides=1,use_bias=False)(x2)
    x2 = tf.keras.layers.BatchNormalization(name='backbone.stem.linear_conv.bn')(x2)
    x2 = tf.keras.layers.Activation("relu")(x2)
    
    #branch1
    x1 = tf.keras.layers.ZeroPadding2D(padding=1)(x1)
    x1 = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=2,use_bias=False,name='backbone.stem.branch1.0.conv')(x1)
    x1 = tf.keras.layers.BatchNormalization(name='backbone.stem.branch1.0.bn')(x1)

    x1 = tf.keras.layers.Conv2D(inc_channels,kernel_size=1,strides=1,use_bias=False,name='backbone.stem.branch1.1.conv')(x1)
    x1 = tf.keras.layers.BatchNormalization(name='backbone.stem.branch1.1.bn')(x1)
    x1 = tf.keras.layers.Activation("relu")(x1)

    x = Lambda( lambda x: tf.concat([x[0],x[1]],axis=3))([x1, x2])
    x = channel_shuffle(x,2)
    return x


def IterativeHead(input, in_channels,prefix=''):
    projects = []
    num_branchs = len(in_channels)
    in_channels = in_channels[::-1]
    print('in channels are these',in_channels)
    x = input[::-1]
    y = []
    last_x = None
    for i, s in enumerate(x):
        if last_x is not None:
            last_x = Lambda( lambda x: tf.compat.v1.image.resize(x[0], size=x[1].shape[-3:-1], method=tf.image.ResizeMethod.BILINEAR,align_corners=True))([last_x,s])
            s = tf.keras.layers.Add()([s,last_x])
        if i != num_branchs - 1:
            #Depthwise seperable convolution
            s1 = tf.keras.layers.ZeroPadding2D(padding=1)(s)
            s1 = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=1,use_bias=False,
            name=prefix+'projects.'+str(i)+'.depthwise_conv.conv')(s1)
            s1 = tf.keras.layers.BatchNormalization(name=prefix+'projects.'+str(i)+'.depthwise_conv.bn')(s1)

            # pointwise convolution
            s1 = tf.keras.layers.Conv2D(in_channels[i + 1],kernel_size=1,use_bias=False,
            name=prefix+'projects.'+str(i)+'.pointwise_conv.conv')(s1)
            s1 = tf.keras.layers.BatchNormalization(name=prefix+'projects.'+str(i)+'.pointwise_conv.bn')(s1)
            s = tf.keras.layers.Activation("relu")(s1)
        else:
            #Depthwise seperable convolution
            s1 = tf.keras.layers.ZeroPadding2D(padding=1)(s)
            s1 = tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=1,use_bias=False,
            name=prefix+'projects.'+str(i)+'.depthwise_conv.conv')(s1)
            s1 = tf.keras.layers.BatchNormalization(name=prefix+'projects.'+str(i)+'.depthwise_conv.bn')(s1)

            # pointwise convolution
            s1 = tf.keras.layers.Conv2D(in_channels[i],kernel_size=1,use_bias=False,
            name=prefix+'projects.'+str(i)+'.pointwise_conv.conv')(s1)
            s1 = tf.keras.layers.BatchNormalization(name=prefix+'projects.'+str(i)+'.pointwise_conv.bn')(s1)
            s = tf.keras.layers.Activation("relu")(s1)

        y.append(s)
        last_x = s

    return y[::-1]


def get_output_list(layers_list, x):
        for layer in layers_list:
            x = layer(x)
        return x 

def LiteHRModule(input_x,
            num_branches,
            num_blocks,
            in_channels,
            reduce_ratio,
            module_type,
            multiscale_output=False,
            with_fuse=True,prefix=''):

    def check_branches(num_branches, in_channels):
        """Check input to avoid ValueError."""
        if num_branches != len(in_channels):
            error_msg = f'NUM_BRANCHES({num_branches}) ' \
                f'!= NUM_INCHANNELS({len(in_channels)})'
            raise ValueError(error_msg)
    
    def fuse_block_1(channels,prefix='',firstindex='',mid_index=''):
        return [tf.keras.layers.Conv2D(channels,kernel_size=1,strides=1,use_bias=False,name=prefix+str(firstindex)+'.'+str(mid_index)+'.0'),#,name=prefix+'0.1.0'
        tf.keras.layers.BatchNormalization(name=prefix+str(firstindex)+'.'+str(mid_index)+'.1'),#name=prefix+'0.1.1'
        tf.keras.layers.UpSampling2D(size=2**(j - i),interpolation='nearest')
        ]
    
    def fuse_block_2(channels,i,j,prefix='',firstindex='',mid_index=''):
        conv_downsamples=[]
        for k in range(i - j):
                if k == i - j - 1:
                    conv_downsamples.append(tf.keras.layers.ZeroPadding2D(padding=1))
                    conv_downsamples.append(tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=2,use_bias=False,name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.0'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.1'))
                    conv_downsamples.append(tf.keras.layers.Conv2D(channels[i],kernel_size=1,strides=1,use_bias=False,name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.2'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.3'))
                
                else:
                    conv_downsamples.append(tf.keras.layers.ZeroPadding2D(padding=1))
                    conv_downsamples.append(tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=2,use_bias=False,name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.0'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.1'))
                    conv_downsamples.append(tf.keras.layers.Conv2D(channels[j],kernel_size=1,strides=1,use_bias=False,name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.2'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name=prefix+str(firstindex)+'.'+str(mid_index)+'.'+str(k)+'.3'))
                    conv_downsamples.append(tf.keras.layers.Activation(tf.keras.activations.relu))
        return conv_downsamples

    
    check_branches(num_branches, in_channels)
    num_out_branches = num_branches if multiscale_output else 1

    # create fuse list
    fuse_layers = []
    for i in range(num_out_branches):
            fuse_layer = []
            block_1=0
            block_2=0
            for j in range(num_branches):
                if j > i:
                    fuse_layer.append(fuse_block_1(in_channels[i],prefix=prefix[:-7]+'fuse_layers.',firstindex=i,mid_index=j))
                    print(block_1)
                    block_1+=1
                elif j == i:
                    fuse_layer.append(None)
                else:
                    fuse_layer.append(fuse_block_2(in_channels,i,j,prefix=prefix[:-7]+'fuse_layersb2.',firstindex=i,mid_index=j))
                    block_2+=1
            fuse_layers.append(fuse_layer)

    #
    if num_branches == 1:
        return [ConditionalChannelWeighting(input_x[0],in_channels,
                    strides=1,
                    reduce_ratio=reduce_ratio,prefix=prefix+'0')]
    if module_type == 'LITE':
        x= input_x
        for i in range(num_blocks):
            x=ConditionalChannelWeighting(x,in_channels,
                strides=1,
                reduce_ratio=reduce_ratio,prefix=prefix+str(i))
    out = x
    if with_fuse:
        out_fuse = []
        for i in range(len(fuse_layers)):
            y = out[0] if i == 0 else get_output_list(fuse_layers[i][0],out[0])
            for j in range(num_branches):
                if i == j:
                        y =tf.keras.layers.Add()([y,out[j]])
                else:
                    y =tf.keras.layers.Add()([y,get_output_list(fuse_layers[i][j],out[j])])
            out_fuse.append(tf.keras.layers.Activation(tf.keras.activations.relu)(y))
        out = out_fuse
    elif not multiscale_output:
        out = [out[0]]
    return out

def LiteHRNet(input_x,extra,in_channels=3,zero_init_residual=False):
    num_stages = extra['num_stages']
    stages_spec = extra['stages_spec']

    num_channels_last = [
            extra['stem']['out_channels'],
        ]

    with_head = extra['with_head']

    def _make_stage(input_x,
                    stages_spec,
                    stage_index,
                    in_channels,
                    multiscale_output=True,prefix=''):
        num_modules = stages_spec['num_modules'][stage_index]
        num_branches = stages_spec['num_branches'][stage_index]
        num_blocks = stages_spec['num_blocks'][stage_index]
        reduce_ratio = stages_spec['reduce_ratios'][stage_index]
        with_fuse = stages_spec['with_fuse'][stage_index]
        module_type = stages_spec['module_type'][stage_index]

        modules = []
        x = input_x
        for i in range(num_modules):
            # multi_scale_output is only used last module
            if not multiscale_output and i == num_modules - 1:
                reset_multiscale_output = False
            else:
                reset_multiscale_output = True
                x = LiteHRModule(x,
                    num_branches,
                    num_blocks,
                    in_channels,
                    reduce_ratio,
                    module_type,
                    multiscale_output=reset_multiscale_output,
                    with_fuse=with_fuse,prefix=prefix+'.'+str(i)+'.layers.')
            
            # in_channels = modules[-1].in_channels

        return x, in_channels
    
    def _make_transition_layer(num_channels_pre_layer,
                               num_channels_cur_layer):
        """Make transition layer."""
        num_branches_cur = len(num_channels_cur_layer)
        num_branches_pre = len(num_channels_pre_layer)

        transition_layers = []
        for i in range(num_branches_cur):
            if i < num_branches_pre:
                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
                    transition_layers.append([tf.keras.layers.ZeroPadding2D(padding=1),
                                            tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=1,use_bias=False,
                                            name='backbone.transition'+str(i)+'.0.0'),
                                            tf.keras.layers.BatchNormalization(name='backbone.transition'+str(i)+'.0.1'),
                                            tf.keras.layers.Conv2D(num_channels_cur_layer[i],kernel_size=1,strides=1,use_bias=False,name='backbone.transition'+str(i)+'.0.2'),
                                            tf.keras.layers.BatchNormalization(name='backbone.transition'+str(i)+'.0.3'),
                                            tf.keras.layers.Activation(tf.keras.activations.relu)])
                else:
                    transition_layers.append(None)
            else:
                conv_downsamples = []
                for j in range(i + 1 - num_branches_pre):
                    in_channels = num_channels_pre_layer[-1]
                    out_channels = num_channels_cur_layer[i] \
                        if j == i - num_branches_pre else in_channels
                    conv_downsamples.append(tf.keras.layers.ZeroPadding2D(padding=1))
                    conv_downsamples.append(tf.keras.layers.DepthwiseConv2D(kernel_size=3,strides=2,use_bias=False,
                    name='backbone.transition'+str(i-1)+'.'+str(i)+'.0.0'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name='backbone.transition'+str(i-1)+'.'+str(i)+'.0.1'))
                    conv_downsamples.append(tf.keras.layers.Conv2D(out_channels,kernel_size=1,strides=1,use_bias=False
                    ,name='backbone.transition'+str(i-1)+'.'+str(i)+'.0.2'))
                    conv_downsamples.append(tf.keras.layers.BatchNormalization(name='backbone.transition'+str(i-1)+'.'+str(i)+'.0.3'))
                    conv_downsamples.append(tf.keras.layers.Activation(tf.keras.activations.relu))
                transition_layers.append(conv_downsamples)

        return transition_layers

    """Forward function."""
    # stem forward
    x = Stem(input_x,extra['stem']['stem_channels'],extra['stem']['out_channels'],extra['stem']['expand_ratio'])
    y_list = [x]
    for i in range(num_stages):
        num_channels = stages_spec['num_channels'][i]
        num_channels = [num_channels[i] for i in range(len(num_channels))]
        x_list = []
        # transition
        transition_layer = _make_transition_layer(num_channels_last, num_channels)
        for j in range(stages_spec['num_branches'][i]):
            if transition_layer[j]:
                if j >= len(y_list):
                    x_list.append(get_output_list(transition_layer[j],y_list[-1]))
                else:
                    x_list.append(get_output_list(transition_layer[j],y_list[j]))
            else:
                x_list.append(y_list[j])
        y_list,num_channels_last = _make_stage(x_list,stages_spec, i, num_channels, multiscale_output=True,prefix='backbone.stage'+str(i))
    x = y_list
    if with_head:
        x = IterativeHead(x,
            in_channels=num_channels_last,prefix='backbone.head_layer.'
        )
    final_out = tf.keras.layers.Conv2D(17,kernel_size=1,strides=1,name='keypoint_head.final_layer')(x[0])
    return final_out


def create_litehrnet():
    extras ={'stem': {'stem_channels': 32, 'out_channels': 32, 'expand_ratio': 1},
     'num_stages': 3, 'stages_spec': {'num_modules': (3, 8, 3), 'num_branches': (2, 3, 4), 'num_blocks': (2, 2, 2),
      'module_type': ('LITE', 'LITE', 'LITE'), 'with_fuse': (True, True, True),
    'reduce_ratios': (8, 8, 8), 'num_channels': ((40, 80), (40, 80, 160), (40, 80, 160, 320))}, 'with_head': True}

    litehrnet_input = tf.keras.Input(shape=(256, 192, 3), name="img")
    outputs = LiteHRNet(litehrnet_input,extras)
    hrnet = tf.keras.Model(inputs=litehrnet_input, outputs=outputs)
    return hrnet

from tensorflow_model_optimization.python.core.quantization.keras import quantize_config
from tensorflow_model_optimization.python.core.quantization.keras import quantizers


class Default8BitOutputQuantizeConfig(quantize_config.QuantizeConfig):
    """QuantizeConfig which only quantizes the output from a layer."""
    def get_weights_and_quantizers(self, layer):
        return []
    
    def get_activations_and_quantizers(self, layer):
        return []

    def set_quantize_weights(self, layer, quantize_weights):
        pass

    def set_quantize_activations(self, layer, quantize_activations):
        pass

    def get_output_quantizers(self, layer):
        return [quantizers.MovingAverageQuantizer(
            num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]

    def get_config(self):
        return {}

class BN8BitOutputQuantizeConfig(quantize_config.QuantizeConfig):
    """QuantizeConfig which only quantizes the output from a layer."""
    def get_weights_and_quantizers(self, layer):
        return []
    
    def get_activations_and_quantizers(self, layer):
        return []

    def set_quantize_weights(self, layer, quantize_weights):
        pass

    def set_quantize_activations(self, layer, quantize_activations):
        pass

    def get_output_quantizers(self, layer):
        return [quantizers.MovingAverageQuantizer(
            num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]

    def get_config(self):
        return {}

import tensorflow_model_optimization as tfmot

def apply_quantization_to_all_except(layer):
    if isinstance(layer, tf.keras.layers.Lambda):
        return layer
    
    if isinstance(layer, tf.keras.layers.Multiply):
        return tfmot.quantization.keras.quantize_annotate_layer(layer,Default8BitOutputQuantizeConfig())

    # if isinstance(layer, tf.keras.layers.BatchNormalization):
    #     if 'fuse_layers' in layer.get_config()['name']:
    #         print(layer,layer.get_config())
    #         # return layer
    #         return tfmot.quantization.keras.quantize_annotate_layer(layer,Default8BitOutputQuantizeConfig())
    #     else:
    #         tfmot.quantization.keras.quantize_annotate_layer(layer)
    return tfmot.quantization.keras.quantize_annotate_layer(layer)

if __name__ == '__main__':
    hrnet = create_litehrnet()
    var = hrnet(tf.random.normal((1,256,192,3)))
    annotated_model = tf.keras.models.clone_model(hrnet,clone_function=apply_quantization_to_all_except)
    # model = tfmot.quantization.keras.quantize_apply(annotated_model)
    quantize_scope = tfmot.quantization.keras.quantize_scope
    with quantize_scope(
        {'DefaultDenseQuantizeConfig': Default8BitOutputQuantizeConfig,
        'CustomLayer': tf.keras.layers.BatchNormalization}):
        # Use `quantize_apply` to actually make the model quantization aware.
        model = tfmot.quantization.keras.quantize_apply(annotated_model)

    model.summary()

only layers in fuse_block 2 are not getting quantized
Screenshots

Additional context
Multiply operation was passed with Default8BitOutputQuantizeConfig to Quantize.
any suggestions and ideas are welcome to make the model completely quantizable

Start with the provided reproduction and the quantization-aware training overview, then trace the fuse_layer handling of the Conv2D, DepthwiseConv2D, and BatchNormalization sequence. Done means the batch-normalization layers in this fuse block are fused as expected and the reproduction no longer reports them as unsupported during QAT.