Hey there,

maybe I am wrong, correct if it was, in `C4_W3_Lab_1_VAE_MNIST`

lab.

So according to the doc, we can use add_loss, when the `loss`

is related to `inputs`

, right?

In the `training loop`

section, the origin code:

```
# compute reconstruction loss
flattened_inputs = tf.reshape(x_batch_train, shape=[-1])
flattened_outputs = tf.reshape(reconstructed, shape=[-1])
loss = bce_loss(flattened_inputs, flattened_outputs) * 784
```

which works for binary `loss`

is actually related to the `inputs`

and final `decoded outputs`

by flattening them.

so, my idea is to move those to `vae_model`

method:

```
def vae_model(encoder, decoder, input_shape):
inputs = tf.keras.layers.Input(shape=input_shape)
mu, sigma, z = encoder(inputs)
reconstructed = decoder(z)
inputs_flatten = tf.keras.layers.Flatten()(inputs)
reconstructed_flatten = tf.keras.layers.Flatten()(reconstructed)
model = tf.keras.Model(inputs=inputs, outputs=reconstructed)
loss = kl_reconstruction_loss(inputs, z, mu, sigma)
model.add_loss(loss)
# NEW: Add the binary_crossentropy loss between inputs and decoded outputs.
loss = tf.keras.losses.binary_crossentropy(inputs_flatten, reconstructed_flatten) * 784
model.add_loss(loss)
return model
```

and then the `training loop`

will be simplified like:

```
with tf.GradientTape() as tape:
# feed a batch to the VAE model
vae(x_batch_train)
# add KLD regularization loss
loss = sum(vae.losses)
# get the gradients and update the weights
grads = tape.gradient(loss, vae.trainable_weights)
optimizer.apply_gradients(zip(grads, vae.trainable_weights))
# compute the loss metric
loss_metric(loss)
```

Will it be a problem comparing the original solution, the training was running in the same as the original one, but I am not sure if I am on the right page?

Thanks