Fixed bug in KL_loss calculation for VAE validation step during training

Author: yasinzaii

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",