Unable to figure out the following ValueError for a Variational Autoencoder implementation

I am trying to fit a variational autoencoder (VAE) with custom loss function using Keras and Tensorflow. It’s throwing a value error. I will provide the traceback. But, for the sake of completeness, let me first provide the definition of the VAE architecture.

def sVAEmodel(params):
    input_dim = params['input_dim']
    latent_dim = 4
    intermediate_dim = 32
    intermediate_dim_2  = 16 # extra added Delete later

    # encoder
    inputs = Input(shape=(input_dim,),name='encoder_input')
    x = Dense(intermediate_dim, activation='relu')(inputs)
    x = Dense(intermediate_dim_2, activation='relu')(x)      
    z_mean = Dense(latent_dim, name='z_mean')(x)
    z_log_var = Dense(latent_dim, name='z_log_var')(x)

    # reparametrization trick to sample latent code (z) from Q(z|x)
    z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

    # instantiate encoder model
    encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')
    #encoder.summary()

    # decoder
    latent_inputs = Input(shape=(latent_dim,), name='z_sampling')
    x = Dense(intermediate_dim_2, activation='relu')(latent_inputs) 
    x = Dense(intermediate_dim, activation='relu')(x)              
    outputs_mean = Dense(input_dim, activation='linear',name='output_mean')(x)
    outputs_var = Dense(input_dim, activation='linear',name='output_var')(x)
    # instantiate decoder model
    decoder = Model(latent_inputs, [outputs_mean,outputs_var], name='decoder')
    #decoder.summary()

    # instantiate vae model
    likelihood_mu,likelihood_log_var = decoder(encoder(inputs)[2])
    vae = Model(inputs, [likelihood_mu,likelihood_log_var], name='vae_mlp')

    #get the loss function
    loss = sVAEloss(z_mean,z_log_var,input_dim)
    #compile
    vae.compile(optimizer= 'adam',loss=loss)

    return vae,encoder,decoder

Where the custom loss function (sVAEloss) is defined as follows:

def sVAEloss(z_mean, z_log_var, input_dim):
    def loss(inputs,outputs):
        """ VAE loss = reconnstruction loss + KL loss"""
        
        log_var = K.variable(0.0)
        
        reconn_loss = 0.5*input_dim*(mse(inputs,outputs)/K.exp(log_var)+log_var+ np.log(2*np.pi))

        #kl loss
        kl_loss = K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        kl_loss *= -0.5

        vae_loss = K.sum(reconn_loss + kl_loss)

        return vae_loss
    return loss

The traceback is provided below:

sigvae,sigvae_encoder,sigvae_decoder = sVAEmodel(params=ae_params)

history = sigvae.fit(train_x, train_x, epochs= 30, batch_size=256).history

---

Traceback (most recent call last):

  File "<ipython-input-69-54e4a7b079f7>", line 1, in <module>
    history = sigvae.fit(train_x, train_x, epochs= 30, batch_size=256).history

  File "/Users/td/anaconda3/envs/py36/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 776, in fit
    shuffle=shuffle)

  File "/Users/td/anaconda3/envs/py36/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 2421, in _standardize_user_data
    exception_prefix='target')

  File "/Users/td/anaconda3/envs/py36/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_utils.py", line 323, in standardize_input_data
    str(len(data)) + ' arrays: ' + str(data)[:200] + '...')

**ValueError: Error when checking model target: the list of Numpy arrays that you are passing to your model is not the size the model expected. Expected to see 2 array(s), but instead got the following list of 1 arrays**: [array([[0.00106576, 0.00203555, 0.02108434, ..., 0.        , 0.        ,
        0.        ],
       [0.00236944, 0.00361116, 0.01054217, ..., 0.        , 0.        ,
        0.        ],
       [0.0...

I checked train_x and it is Numpy.ndarray type. I tried np.array(train_x) but that didn’t work.

Note that the decoder has two outputs, like the encoder as I am trying to learn p(x|z) ~ N(mu, sigma**2I) instead of p(x|z) ~ N(mu, I) and I want to calculate the probability of the reconstructed input as well.

Sorry for the post being so long. But I thought I should provide the complete picture of what I am trying to accomplish. Thanks for your help. Your suggestions will be greatly appreciated as I am stuck right now.