NotImplementedError: Cannot convert a symbolic Tensor (truediv_4577:0) to a numpy array.

Hi,
I am trying to train variational autoencoder with reparamaterization trick. In similar implementation on colab, I got error as
_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'latent_sigma/Identity:0' shape=(None, 2) dtype=float32>, <tf.Tensor 'latent_mu/Identity:0' shape=(None, 2) dtype=float32>]

Here is the colab link: https://colab.research.google.com/drive/1_TjoHxDMC3QPQxO9Un1LvMQFHpVAMlMd

Here is my code:

# Definition
i       = Input(shape=input_shape, name='encoder_input')
cx      = Conv2D(filters=8, kernel_size=3, strides=2, padding='same', activation='relu')(i)
cx      = BatchNormalization()(cx)
cx      = Conv2D(filters=16, kernel_size=3, strides=2, padding='same', activation='relu')(cx)
cx      = BatchNormalization()(cx)
x       = Flatten()(cx)
x       = Dense(20, activation='relu')(x)
x       = BatchNormalization()(x)
mu      = Dense(latent_dim, name='latent_mu')(x)
sigma   = Dense(latent_dim, name='latent_sigma')(x)

def sample_z(args):
    mu, sigma = args
    batch     = k.shape(mu)[0]
    dim       = k.shape(mu)[1]
    eps       = k.random_normal(shape=(batch, dim))
    return mu + k.exp(0.5 * sigma) * eps

z = Lambda(sample_z, output_shape=(latent_dim, ), name='z')([mu, sigma])

encoder = Model(i, [mu, sigma, z], name='encoder')
encoder.summary()

conv_shape = K.int_shape(cx)

# Definition
d_i   = Input(shape=(latent_dim, ), name='decoder_input')
x     = Dense(conv_shape[1] * conv_shape[2] * conv_shape[3], activation='relu')(d_i)
x     = BatchNormalization()(x)
x     = Reshape((conv_shape[1], conv_shape[2], conv_shape[3]))(x)
cx    = Conv2DTranspose(filters=16, kernel_size=3, strides=2, padding='same', activation='relu')(x)
cx    = BatchNormalization()(cx)
cx    = Conv2DTranspose(filters=8, kernel_size=3, strides=2, padding='same',  activation='relu')(cx)
cx    = BatchNormalization()(cx)
o     = Conv2DTranspose(filters=num_channels, kernel_size=3, activation='sigmoid', padding='same', name='decoder_output')(cx)

# Instantiate decoder
decoder = Model(d_i, o, name='decoder')
decoder.summary()

# Instantiate VAE
vae_outputs = decoder(encoder(i)[2])
vae         = Model(i, vae_outputs, name='vae')
vae.summary()

opt = tf.keras.optimizers.Adam()


def kl_reconstruction_loss(true, pred):
    # Reconstruction loss
    reconstruction_loss = binary_crossentropy(K.flatten(true), K.flatten(pred))# * img_width * img_height
    kl_loss = tf.constant(0.5) * k.sum(1 + sigma - k.square(mu) - k.exp(sigma), axis=1)
#     print(type(kl_loss))
    return K.mean(reconstruction_loss + kl_loss)

def step(X, y):
    # keep track of our gradients
    with tf.GradientTape() as tape:
        # make a prediction using the model and then calculate the
        # loss
        pred = vae(X)
#         loss = tf.reduce_mean()
        loss = kl_reconstruction_loss(y, pred)
    # calculate the gradients using our tape and then update the
    # model weights
    grads = tape.gradient(loss, vae.trainable_variables)
    opt.apply_gradients(zip(grads, vae.trainable_variables))

EPOCHS = 25
BS = 64
numUpdates = int(input_train.shape[0] / BS)

for epoch in range(0, EPOCHS):
    # loop over the data in batch size increments
    for i in range(0, numUpdates):
        start = i * BS
        end = start + BS
        step(input_train[start:end], input_train[start:end])
        print("Done for batch: ", i)        
    # show timing information for the epoch

vae.compile(optimizer=opt, loss=kl_reconstruction_loss, metrics=["acc"], experimental_run_tf_function=False)

(loss, acc) = vae.evaluate(input_test, input_test)

I got error while running

(loss, acc) = vae.evaluate(input_test, input_test)
And the error is:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-219-014b7d06d698> in <module>
----> 1 (loss, acc) = vae.evaluate(input_test, input_test)

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py in evaluate(self, x, y, batch_size, verbose, sample_weight, steps, callbacks, max_queue_size, workers, use_multiprocessing)
    928         max_queue_size=max_queue_size,
    929         workers=workers,
--> 930         use_multiprocessing=use_multiprocessing)
    931 
    932   def predict(self,

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py in evaluate(self, model, x, y, batch_size, verbose, sample_weight, steps, callbacks, **kwargs)
    818         verbose=verbose,
    819         workers=0,
--> 820         callbacks=callbacks)
    821 
    822   def predict(self,

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py in model_iteration(model, data, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, validation_freq, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch, mode, batch_size, steps_name, **kwargs)
    293       batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode)
    294       callbacks._call_batch_hook(mode, 'end', step, batch_logs)
--> 295       progbar.on_batch_end(step, batch_logs)
    296       step += 1
    297 

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/keras/callbacks.py in on_batch_end(self, batch, logs)
    779     # will be handled by on_epoch_end.
    780     if self.verbose and (self.target is None or self.seen < self.target):
--> 781       self.progbar.update(self.seen, self.log_values)
    782 
    783   def on_epoch_end(self, epoch, logs=None):

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in update(self, current, values)
    557         info += ' - %s:' % k
    558         if isinstance(self._values[k], list):
--> 559           avg = np.mean(self._values[k][0] / max(1, self._values[k][1]))
    560           if abs(avg) > 1e-3:
    561             info += ' %.4f' % avg

<__array_function__ internals> in mean(*args, **kwargs)

~/Documents/env_tf/lib/python3.7/site-packages/numpy/core/fromnumeric.py in mean(a, axis, dtype, out, keepdims)
   3333 
   3334     return _methods._mean(a, axis=axis, dtype=dtype,
-> 3335                           out=out, **kwargs)
   3336 
   3337 

~/Documents/env_tf/lib/python3.7/site-packages/numpy/core/_methods.py in _mean(a, axis, dtype, out, keepdims)
    133 
    134 def _mean(a, axis=None, dtype=None, out=None, keepdims=False):
--> 135     arr = asanyarray(a)
    136 
    137     is_float16_result = False

~/Documents/env_tf/lib/python3.7/site-packages/numpy/core/_asarray.py in asanyarray(a, dtype, order)
    136 
    137     """
--> 138     return array(a, dtype, copy=False, order=order, subok=True)
    139 
    140 

~/Documents/env_tf/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in __array__(self)
    726   def __array__(self):
    727     raise NotImplementedError("Cannot convert a symbolic Tensor ({}) to a numpy"
--> 728                               " array.".format(self.name))
    729 
    730   def __len__(self):

NotImplementedError: Cannot convert a symbolic Tensor (truediv_4577:0) to a numpy array.

Start with the provided notebook/code and reproduce the failure at vae.evaluate(input_test, input_test). Trace the symbolic tensor passed into the progress bar's NumPy mean. Done means evaluation completes without the NotImplementedError and returns loss and accuracy.