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 <[email protected]>
Co-authored-by: Can Zhao <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]>

Author: 3 people

generation/maisi/maisi_train_vae_tutorial.ipynb

@@ -692,7 +692,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "id": "4c251a32-390f-46dd-a613-75b12a7884c1",
    "metadata": {
     "scrolled": true
@@ -850,7 +850,7 @@
     "            with torch.no_grad():\n",
     "                with autocast(\"cuda\", enabled=args.amp):\n",
     "                    images = batch[\"image\"]\n",
-    "                    reconstruction, _, _ = dynamic_infer(val_inferer, autoencoder, images)\n",
+    "                    reconstruction, z_mu, z_sigma = dynamic_infer(val_inferer, autoencoder, images)\n",
     "                    reconstruction = reconstruction.to(device)\n",
     "                    val_epoch_losses[\"recons_loss\"] += intensity_loss(reconstruction, images.to(device)).item()\n",
     "                    val_epoch_losses[\"kl_loss\"] += KL_loss(z_mu, z_sigma).item()\n",