Как написать замаскированную потерю MSE в Keras?

Я пытался написать замаскированную потерю MSE:

def mae_loss_masked(mask):
    def loss_fn(y_true, y_pred):
        abs_vec = tf.multiply(tf.abs(y_pred-y_true), mask)
        loss = tf.reduce_mean(abs_vec)
        return loss
    return loss_fn

Моя модель:

def MobileNet_v1():
    # MobileNet with dense layer on top

    # Keras 2.1.6
    mobilenet = MobileNet(input_shape=(config.IMAGE_H, config.IMAGE_W, config.N_CHANNELS),
                          alpha=1.0,
                          depth_multiplier=1,
                          include_top=False,
                          weights='imagenet'
                          )

    x = Flatten()(mobilenet.output)
    x = Dropout(0.5)(x)
    x = Dense(config.N_LANDMARKS * 2, activation='linear')(x)

    # -------------------------------------------------------

    model = Model(inputs=mobilenet.input, outputs=x)
    optimizer = Adadelta()
    model.compile(optimizer=optimizer, loss=mae_loss_masked)

    model.summary()
    import sys
    sys.exit()

    return model

Но выдает ошибку: TypeError: mae_loss_masked() takes 1 positional argument but 2 were given

Также вопрос, как в этом случае должен выглядеть вывод пакетного генератора.