TensorFlow 2.4 behaviour with custom loss functions: Cannot convert a symbolic Keras input/output to a numpy array

I'm learning TensorFlow pretty much from scratch, and finding it hard to replicate examples on latest versions of TF.

In a typical VAE example of MNIST I have a segment of code like:

(...)
def vae_loss(y_true, y_pred):
    kl_loss = vae_kl_loss(y_true, y_pred)
    rc_loss = vae_rc_loss(y_true, y_pred)
    kl_weight_const = 1e-2
    return kl_weight_const*kl_loss + rc_loss

vae.compile(
    loss = [vae_loss],
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    metrics=[vae_kl_loss,vae_rc_loss]
)

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.astype('float32')/255
x_test = x_test.astype('float32')/255
x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))

vae.fit(x_train, x_train, shuffle=True, epochs=nb_epoch, batch_size = batch_size, validation_data=(x_test, x_test))

With TensorFlow 2.4 and no special config, receive the error:

Cannot convert a symbolic Keras input/output to a numpy array. This error may indicate that you're trying to pass a symbolic value to a NumPy call, which is not supported. Or, you may be trying to pass Keras symbolic inputs/outputs to a TF API that does not register dispatching, preventing Keras from automatically converting the API call to a lambda layer in the Functional Model.

Through some search, I found references that this could be a symptom of eager of execution in TF2 and indeed with

tf.compat.v1.disable_eager_execution()

The code runs. Now I'm trying to learn what would be the way to do this example with eager execution. Interestingly enough on TensorFlow 2.3 with

tf.config.run_functions_eagerly(True)

it runs, but not with TensorFlow 2.4. In another similar example I have where x_train is loaded with tensorflow_datasets, but otherwise same TF version and loss function, it also works correctly without any config settings. I'd assume that the behaviour is because the input is now a tf.data.Dataset rather than numpy array ?

Thanks in advance for any help you may provide explaining how I can my example working with eager mode.

Thanks