Различная конфигурация слоев для фазы поезда и теста в керасе

Question

ребята, я пытаюсь использовать другой путь для фазы поезда и теста в моем NN, я создал функцию для моей модели с логическим значением, которая указывает фазу поезда или теста, но при загрузке модели с использованием кератов я получил сеть в конфигурации обучения, а нетестовое задание.это означает, что я получаю 3 результата вместо 1.

   net = load_model('output/weights.100-0.860.h5',custom_objects={'train':False})

Это то, где меня вдохновили: https://github.com/aitorzip/Keras-ICNet

    def build_bn(height, width, n_classes, weights_path=None, train=False):
    inp = Input(shape=(height, width, 4))
        .....#some codes


        if train:
        aux_1 = Conv2D(n_classes, 1, activation='softmax', name='sub4_out')(aux_1)
        aux_2 = Conv2D(n_classes, 1, activation='softmax', name='sub24_out')(aux_2)
        model = Model(inputs=inp, outputs=[out, aux_2, aux_1])
        # model.summary()

    else:
        model = Model(inputs=inp, outputs=out)

    if weights_path is not None:
        model.load_weights(weights_path, by_name=True)


    return model