Fixed bug in KL_loss calculation for VAE validation step during training (#2000)

Fixes # .

1. Bug in KL_loss calculation for VAE validation step during training

### Description
The KL_loss is included in the Validation loss calculation during
VAE-training validation step. However, the correct z_mu, z_sigma are not
passed. Either this val_epoch_losses["kl_loss"] should always be zero if
that was the intension or the correct z_mu, z_sigma should be passed.

### Checks
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Avoid including large-size files in the PR.
- [x] Clean up long text outputs from code cells in the notebook.
- [x] For security purposes, please check the contents and remove any
sensitive info such as user names and private key.

---------

Signed-off-by: Muhammad Nabi Yasinzai <muhammad.nabi.yasinzai@gmail.com>
Co-authored-by: Can Zhao <69829124+Can-Zhao@users.noreply.github.com>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>