Change training and testing behavior of custom Layer Keras

I am trying to change my Keras model's behavior during training and testing.

To be more precise, I want to simply evaluate the predictions during training, and add another Lambda layer (as a kind of postprocessing) for testing.

I've found a solution Here, where K.functionreceives the specified K.learning_phase() and returns an output. As I understand, using K.in_test_phase() or K.in_training_phase() will return either the first or the second parameter (as described in the docs), based on the training parameter passed.

I'm running on TF 2.0 as Backend, so eager execution is enabled by default. That being said, passing K.learning_phase() results in an error, which has been described see here. Hence, I'm using tensorflow.python.keras.symbolic_learning_phase(), which seems to work. I'm currently able to get the output with K.function(), but my aim is to perform model.fit() in order to train my model, and later call model.evaluate() (I'm using the Keras' functional API).

How would the proper way to train and test my model, based on the learning flag be?

Currently my MWE is:

def build_model(images, training=None):

    input_layer = Input(shape=(256,256,3), dtype="float32", batch_size=80)

    ...
    #performing some factor disentanglement here


    angle_pred = UpSampling2D(size=(8, 8), interpolation='bilinear')(angle)
    radius_pred = UpSampling2D(size=(8, 8), interpolation='bilinear')(radius)

    angle_radius_stack = tf.stack([angle_pred, radius_pred], 0)
    hough_voting = Lambda(hough_vote)((angle_radius_stack, images))
    train_test = K.in_test_phase(hough_voting, angle_radius_stack, training = training)

    angle_v, radius_v = tf.unstack(train_test)

    model = Model(inputs=input_layer, outputs=[angle_v, radius_v])
    adam = Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, amsgrad=False)
    model.compile(optimizer=adam, loss='mean_absolute_error', metrics=['accuracy'], run_eagerly=True) 
    return model

And then using training:

def train_model(model, patches, radii, angles):

    fun = K.function([model.layers[0].input, B.symbolic_learning_phase()], [model.layers[-1].output])
    print(fun([patches[0:80,:,:,:], True]))

    model.fit(patches, [radii, angles], batch_size=80, epochs=1)