title: Call method not implemented runtime error, when creating VAE in Keras. Model Subclassing
body: <p>I'm running the VAE Keras code sample from here: <a href="https://keras.io/examples/generative/vae/" rel="nofollow noreferrer">https://keras.io/examples/generative/vae/</a> and it works OK. But when I modify it to use a validation_split = 0.2 in vae.fit() method I'm getting a runtime error: NotImplementedError: When subclassing the <code>Model</code> class, you should implement a <code>call</code> method.</p>
  <p>The code uses Model subclassing. My question is: Why it works well without using validation_split parameter on fit()? and what I have to do to get the VAE working with a validation_split?</p>
  <p>I'm already tried to remove the model names as suggested here: <a href="https://stackoverflow.com/questions/57201943/during-creating-vae-model-throws-exception-you-should-implement-a-call-method">During creating VAE model throws exception &quot;you should implement a `call` method.&quot;</a> but the problem persists.</p>
  <pre><code>decoder =keras.Model(latent_inputs, decoder_outputs)
  encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z])
  </code></pre>
  <p>Here is the Model subclass:</p>
  <pre><code>class VAE(keras.Model):
  def __init__(self, encoder, decoder, **kwargs):
      super(VAE, self).__init__(**kwargs)
      self.encoder = encoder
      self.decoder = decoder
  
  def train_step(self, data):
      if isinstance(data, tuple):
          data = data[0]
      with tf.GradientTape() as tape:
          z_mean, z_log_var, z = encoder(data)
          reconstruction = decoder(z)
          reconstruction_loss = tf.reduce_mean(
              keras.losses.binary_crossentropy(data, reconstruction)
          )
          reconstruction_loss *= 28 * 28
          kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
          kl_loss = tf.reduce_mean(kl_loss)
          kl_loss *= -0.5
          total_loss = reconstruction_loss + kl_loss
      grads = tape.gradient(total_loss, self.trainable_weights)
      self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
      return {
          &quot;loss&quot;: total_loss,
          &quot;reconstruction_loss&quot;: reconstruction_loss,
          &quot;kl_loss&quot;: kl_loss,
      }
  </code></pre>
  <p>Sampling layer:</p>
  <pre><code>class Sampling(layers.Layer):
  &quot;&quot;&quot;Uses (z_mean, z_log_var) to sample z, the vector encoding a digit.&quot;&quot;&quot;
  
  def call(self, inputs):
      z_mean, z_log_var = inputs
      batch = tf.shape(z_mean)[0]
      dim = tf.shape(z_mean)[1]
      epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
      return z_mean + tf.exp(0.5 * z_log_var) * epsilon
  </code></pre>
  <p>Encoder definition:</p>
  <pre><code>latent_dim = 2
  
  encoder_inputs = keras.Input(shape=(28, 28, 1))
  x = layers.Conv2D(32, 3, activation=&quot;relu&quot;, strides=2, padding=&quot;same&quot;)(encoder_inputs)
  x = layers.Conv2D(64, 3, activation=&quot;relu&quot;, strides=2, padding=&quot;same&quot;)(x)
  x = layers.Flatten()(x)
  x = layers.Dense(16, activation=&quot;relu&quot;)(x)
  z_mean = layers.Dense(latent_dim, name=&quot;z_mean&quot;)(x)
  z_log_var = layers.Dense(latent_dim, name=&quot;z_log_var&quot;)(x)
  z = Sampling()([z_mean, z_log_var])
  encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name=&quot;encoder&quot;)
  encoder.summary()
  </code></pre>
  <p>Decoder definition:</p>
  <pre><code>latent_inputs = keras.Input(shape=(latent_dim,))
  x = layers.Dense(7 * 7 * 64, activation=&quot;relu&quot;)(latent_inputs)
  x = layers.Reshape((7, 7, 64))(x)
  x = layers.Conv2DTranspose(64, 3, activation=&quot;relu&quot;, strides=2, padding=&quot;same&quot;)(x)
  x = layers.Conv2DTranspose(32, 3, activation=&quot;relu&quot;, strides=2, padding=&quot;same&quot;)(x)
  decoder_outputs = layers.Conv2DTranspose(1, 3, activation=&quot;sigmoid&quot;, padding=&quot;same&quot;)(x)
  decoder = keras.Model(latent_inputs, decoder_outputs, name=&quot;decoder&quot;)
  decoder.summary()
  </code></pre>
  <p>Train VAE</p>
  <pre><code>(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
  mnist_digits = np.concatenate([x_train, x_test], axis=0)
  mnist_digits = np.expand_dims(mnist_digits, -1).astype(&quot;float32&quot;) / 255
  
  vae = VAE(encoder, decoder)
  vae.compile(optimizer=keras.optimizers.Adam())
  vae.fit(mnist_digits, epochs=30, batch_size=128, validation_split=0.2)
  </code></pre>
  <p>Thanks in advance!</p>