I have trained InceptionResnetV2 using CBAM. Heres the cbam i am using

class ChannelAttention(Layer):
    def __init__(self, ratio=8):
        super(ChannelAttention, self).__init__()
        self.ratio = ratio

    def build(self, input_shape):
        self.shared_layer_one = Dense(input_shape[-1] // self.ratio,
                                      activation='relu',
                                      kernel_initializer='he_normal',
                                      use_bias=True,
                                      bias_initializer='zeros')
        self.shared_layer_two = Dense(input_shape[-1],
                                      kernel_initializer='he_normal',
                                      use_bias=True,
                                      bias_initializer='zeros')

    def call(self, input_tensor):
        avg_pool = GlobalAveragePooling2D()(input_tensor)
        avg_pool = Reshape((1, 1, avg_pool.shape[-1]))(avg_pool)
        assert avg_pool.shape[1:] == (1, 1, input_tensor.shape[-1])
        avg_pool = self.shared_layer_one(avg_pool)
        avg_pool = self.shared_layer_two(avg_pool)

        max_pool = GlobalMaxPooling2D()(input_tensor)
        max_pool = Reshape((1, 1, max_pool.shape[-1]))(max_pool)
        assert max_pool.shape[1:] == (1, 1, input_tensor.shape[-1])
        max_pool = self.shared_layer_one(max_pool)
        max_pool = self.shared_layer_two(max_pool)

        cbam_feature = Add()([avg_pool, max_pool])
        cbam_feature = Activation('sigmoid')(cbam_feature)

        return Multiply()([input_tensor, cbam_feature])

class SpatialAttention(Layer):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        self.conv2d = Conv2D(1, kernel_size, padding='same', activation='sigmoid', kernel_initializer='he_normal')

    def call(self, input_tensor):
        avg_pool = tf.reduce_mean(input_tensor, axis=-1, keepdims=True)
        max_pool = tf.reduce_max(input_tensor, axis=-1, keepdims=True)
        concat = tf.concat([avg_pool, max_pool], axis=-1)
        cbam_feature = self.conv2d(concat)  
        return Multiply()([input_tensor, cbam_feature])

class CBAMBlock(Layer):
    def __init__(self, ratio=8, kernel_size=7):
        super(CBAMBlock, self).__init__()
        self.channel_attention = ChannelAttention(ratio)
        self.spatial_attention = SpatialAttention(kernel_size)

    def call(self, input_tensor):
        x = self.channel_attention(input_tensor)
        x = self.spatial_attention(x)
        return x

There was no problem in training the model. The problem that I am facing now is that after I saved the model as “inceptionResnetV2.h5” and tried to load by passing a custom object like this

custom = {
    'ChannelAttention': ChannelAttention,
    'SpatialAttention': SpatialAttention,
    'CBAMBlock': CBAMBlock
}

model = load_model('inceptionResnetV2.h5', custom_objects=custom)

now I am getting the ValueError.

ValueError: Unknown layer: 'CustomScaleLayer'. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

What should I do to load my model now?