Почему моя модель не проходит go через пользовательский слой в Tensorflow Keras?

Question

Я пытаюсь написать свою собственную функцию потерь через пользовательские слои. В функции потерь рассчитывается отклонение весов между двумя разными слоями. Вот моя функция потерь и как ее вызвать:

class WEIGHTS_LOSS(Layer):
    def __init__(self, **kwargs):
        super(WEIGHTS_LOSS, self).__init__(**kwargs)
        self.b = tf.Variable(initial_value=tf.zeros((1,)), trainable=True)
        self.a = tf.Variable(initial_value=tf.ones((1,)), trainable=True)

    def call(self, inputs, **kwargs):
        target_conv, source_conv, target_features, source_features = \
            inputs
        target_weights = K.reshape(target_conv[0], (-1, 1))
        source_weights = K.reshape(source_conv[0], (-1, 1))
        target_bias = K.reshape(target_conv[1], (-1, 1))
        source_bias = K.reshape(source_conv[1], (-1, 1))
        target_conv_data = K.concatenate([target_weights, target_bias], axis=0)
        source_conv_data = K.concatenate([source_weights, source_bias], axis=0)
        weights_loss = K.exp(K.sum(K.square(self.a * source_conv_data + self.b - target_conv_data)))-1
        self.add_loss(weights_loss, inputs=True)
        return inputs[2], inputs[3]

    def get_config(self, **kwargs):
        super(WEIGHTS_LOSS, self).get_config(**kwargs)

if __name__ == '__main__':
    custom_layer = WEIGHTS_LOSS()
    input1 = Input((224, 224, 1))
    input2 = Input((224, 224, 1))
    conv1 = Conv2D(16, 3)
    conv2 = Conv2D(16, 3)
    x1 = conv1(input1)
    x2 = conv2(input2)
    outputs = custom_layer([conv1.weights, conv2.weights, x1, x2])
    model = Model([input1, input2], outputs)
    model.summary()

Однако, когда я запускаю model.summary(), пользовательский слой не включается в мою модель.

Model: "model"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 224, 224, 1) 0                                            
__________________________________________________________________________________________________
input_2 (InputLayer)            [(None, 224, 224, 1) 0                                            
__________________________________________________________________________________________________
conv2d (Conv2D)                 (None, 222, 222, 16) 160         input_1[0][0]                    
__________________________________________________________________________________________________
conv2d_1 (Conv2D)               (None, 222, 222, 16) 160         input_2[0][0]                    
==================================================================================================
Total params: 320
Trainable params: 320
Non-trainable params: 0
__________________________________________________________________________________________________

Process finished with exit code 0

Я действительно запутался в результате, поскольку другой пользовательский слой, который я написал, был успешно пройден.
В чем, возможно, причина этого и как я могу решить эту проблему?