Tensorflow cvae error when not using tf.compat.v1.disable_v2_behavior

Question

I am running a toy Variational autoencoder (VAE) as described here https://blog.keras.io/building-autoencoders-in-keras.html, just using mse reconstruction loss.

act1 = 'selu'
n1 = tf.keras.Input(shape = (input_1+ input_2,) )
n2 =  tf.keras.layers.Dense(units=64, activation=act1)(n1)
n4 =  tf.keras.layers.Dense(units=16, activation=act1)(n2)

def sample_z(args):
    z_mean, z_log_var = args
    batch_size = tf.shape(z_mean)[0]
    dim = tf.shape(z_mean)[1]
    eps = tf.random.normal(shape=(batch_size, dim), mean=0.0, stddev=1.0)
    return z_mean + tf.math.exp(z_log_var * 0.5) * eps

en6_mean =  tf.keras.layers.Dense(units=latent_dim, activation=None, name='z_mean')(n4)
en6_log_var =  tf.keras.layers.Dense(units=latent_dim, activation=None, name='z_log_var')(n4)
en6_z = tf.keras.layers.Lambda(sample_z, output_shape = (latent_dim, ), name='z')([en6_mean, en6_log_var])
encoder_model = tf.keras.Model(inputs=en1, outputs=[en6_mean, en6_log_var, en6_z], name = 'encoder' )
encoder_model.summary()


d1 =  tf.keras.Input(shape = (latent_dim + input_2,))  
d2 =  tf.keras.layers.Dense(units = 16, activation = act1)(d1)
d4 =  tf.keras.layers.Dense(units = 64, activation = act1)(d2)
d5 =  tf.keras.layers.Dense(units = input_1+ input_2, activation = None)(d4)
decoder_model = tf.keras.Model(inputs=d1, outputs=d5, name = 'decoder' )
decoder_model.summary()                           


temp14 = encoder_model(n1)[2] 
reconstruction = decoder_model( tf.concat([encoder_model(n1)[2], input_2], axis=1)  ) 
cvae_model = tf.keras.Model(inputs = n1, 
                            outputs = reconstruction, 
                            name = 'cvae' )      
cvae_model.summary()
   
    
temp11 = tf.concat([input_1, input_2], axis=1)
recon_loss1 = tf.keras.losses.MeanSquaredError()(temp11, reconstruction)    
kl_loss1 = tf.math.reduce_sum(-0.5 * (1 + en6_log_var - tf.math.exp(en6_log_var) - tf.math.square(en6_mean) ))

total_cvae_loss = tf.math.reduce_mean(recon_loss1 + kl_loss1)
cvae_model.add_loss(total_cvae_loss)
cvae_model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate=lr1) )


cvae_hist = cvae_model.fit(x = tf.concat([input_1, input_2], axis=1), 
                           batch_size = num_batch, epochs = num_epochs, 
                           steps_per_epoch = steps_per_epoch1, verbose = 1, )

However, I am getting error in d1.

InvalidArgumentError:  ConcatOp : Dimensions of inputs should match: shape[0] = [8000,2] vs. shape[1] = [4234,1]
Errors may have originated from an input operation.
Input Source operations connected to node cvae/tf.concat/concat:
Function call stack:
train_function

The error goes away if I use tensorflow.compat.v1.disable_v2_behavior() Any idea why default TF2 behavior causes error?