models/GANs/notebooks/VAE_GAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variational Autoencoder Generative Adversarial Network (VAE-GAN)\n",
    "```\n",
    "This is a VAE-GAN implementation in PyTorch.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# necessary imports\n",
    "# add ../src to sys.path\n",
    "import sys\n",
    "sys.path.append('../src')\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "import VAE_GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load MNIST data\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
    "\n",
    "# datasets\n",
    "trainset = torchvision.datasets.MNIST(root='../../data', train=True, download=True, transform=transform)\n",
    "testset = torchvision.datasets.MNIST(root='../../data', train=False, download=True, transform=transform)\n",
    "\n",
    "\n",
    "# data loaders\n",
    "batch_size = 128\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAdMAAAHqCAYAAABfi6TIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAaDklEQVR4nO3de5CWZf0/8GtXMoUyilZlRaSpZUqUycZxQYIxzZTKTZSjh8FTRo0lMx2cUswEmRwndcqxQccBz2yuoZBRIzOpyEGdxGaoHEgHRUTxkJpDeNr9/tHoL392Xc/yfNh99llerxn/ed7c13PtsrdvbnY/XA1dXV1dCQCoWmOtNwAA9U6ZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkCQMq1DGzduTNOnT0/Dhg1LAwcOTJ/97GfTpZdemrZv317rrQHd9MYbb6QLLrggNTc3p7333ju1trame++9t9bbokoN/m3e+rJ58+Y0evTo9LGPfSzNmjUrfeITn0hr1qxJixYtSm1tbenuu++u9RaBbpgxY0bq6OhIs2fPTi0tLWnRokXpkUceSX/605/SF7/4xVpvj52kTOvM/Pnz04UXXpjWr1+fRo0a9d7rM2fOTDfddFN6+eWX08c//vEa7hCo5OGHH06tra3piiuuSD/4wQ9SSint2LEjHXLIIWnfffdNq1evrvEO2Vn+mrfOvPbaaymllPbbb7/3vT506NDU2NiY9txzz1psC9gJHR0daY899kjnnnvue6/ttdde6eyzz05r1qxJmzdvruHuqIYyrTNHHXVUSimls88+Oz322GNp8+bNqb29Pf36179O3/ve99KgQYNqu0GgonXr1qWRI0emffbZ532vH3HEESmllB577LEa7IqIAbXeADvn+OOPT3Pnzk3z589PS5cufe/1Cy+8MM2bN6+GOwO6a+vWrWno0KEfeP3d15599tne3hJByrQOjRgxIk2YMCGdfPLJaciQIemee+5J8+fPT/vvv38677zzar09oIJ///vf6cMf/vAHXt9rr73ey6kvyrTOLF68OJ177rlpw4YNadiwYSmllE466aTU2dmZLrjggjRjxow0ZMiQGu8SKNl7773TG2+88YHXd+zY8V5OffE90zpz7bXXpsMOO+y9In1XW1tb2r59e1q3bl2NdgZ019ChQ9PWrVs/8Pq7rzU3N/f2lghSpnXm+eefT++8884HXn/rrbdSSim9/fbbvb0lYCd9/vOfTxs2bHjvp/Pf9dBDD72XU1+UaZ0ZOXJkWrduXdqwYcP7Xr/99ttTY2NjGj16dI12BnTX5MmT0zvvvJOuu+66915744030sKFC1Nra2s68MADa7g7quF7pnXmhz/8YVq+fHkaP358Ou+889KQIUPS7373u7R8+fJ0zjnn+OshqAOtra1pypQp6cc//nHatm1b+sxnPpNuvPHGtGnTpnTDDTfUentUwb+AVIcefvjhdMkll6R169all156KX3qU59KM2fOTD/60Y/SgAH+fAT1YMeOHWnOnDnplltuSf/85z/T6NGj09y5c9Nxxx1X661RBWUKAEG+ZwoAQcoUAIKUKQAEKVMACFKmABCkTAEgSJkCQFC3J/wbGhp6ch/QL/W1MW73Mey87tzHnkwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACCBtR6AwD0nIEDB2azjRs3Fq9tbm7OZqtXr85m48aNq7yxfsaTKQAEKVMACFKmABCkTAEgSJkCQJAyBYAgozEAde6jH/1oNlu4cGE222+//YrrdnZ2VpXtjjyZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQZM60jxgzZkwxHz58eDY7//zzs9mRRx5ZXLc0K9bYmP+zVqUZsz322KOYA7vOSSedlM1OPPHEqtf9y1/+ks2uvPLKqtftjzyZAkCQMgWAIGUKAEHKFACClCkABClTAAgyGtOLSuMvixcvLl574IEHZrPIMUnVHqPk+CW6q7m5uZh3dXVls61bt+7q7dSlu+66q5gfc8wxPfK+N998czZbsmRJj7xnvfJkCgBByhQAgpQpAAQpUwAIUqYAEKRMASCooav0c+n//QsbGnp6L3WhNKKSUnnEpXSCS6VRky1btmSz1atXZ7NKv2+l3/7StVdddVVx3bVr1xbz3UU3b69e4z7uu0rjL21tbcVrN23alM3+8Ic/ZLNzzjmnuO4LL7yQzVpaWrLZ9u3bi+vWm+7cx55MASBImQJAkDIFgCBlCgBByhQAgpQpAAQpUwAIMme6k1atWlXMjzjiiGzW2Jj/s0ulOdPx48dnMzOdfZc5U/7bpEmTstlNN92UzUrznimldPnll2ezBQsWZLOnnnqquO6wYcOy2ezZs7PZr371q+K69cacKQD0AmUKAEHKFACClCkABClTAAhSpgAQNKDWG+iLxowZk81Kx6ilVB5x6ejoyGbTpk2rvLFeVjpurvQ5Siml9vb2bFb6MfPS5yilvvl5gneVRl9SSmnhwoXZbO+9985mlUZNSuMvPWXq1KnZrL+NxnSHJ1MACFKmABCkTAEgSJkCQJAyBYAgZQoAQcoUAILMmf4PpaOFKh2VVsr72nFclZRmSW+77bbitaWPtT99jtj9fP3rX89mpWPUUirPkn7nO9/JZrfffnvljVXh0UcfLealI9h4P0+mABCkTAEgSJkCQJAyBYAgZQoAQcoUAIKMxvwPDz30UDaLHAFWunbz5s3Fa++8885stnbt2mxWOkYtpZQWL16czcaOHZvNKo2wNDQ0ZLPSfqdPn15cF3pDc3NzNrvkkkuy2YAB5f+lXnfddVVlPeWBBx4o5t/4xjeyWWOjZ7H/5rMBAEHKFACClCkABClTAAhSpgAQpEwBIKihq5vHdJRGHXYnV1xxRTEvnThT+lHySqfRPPvss9lszZo12azSaMwRRxyRzSL7LY2/zJgxI5s988wzxXXrTV87Bcd9/B9Dhgwp5uvXr89m++67bzb77W9/W1x3ypQp5Y31shUrVhTzL33pS9ls9erV2Wz8+PFV76kv6s597MkUAIKUKQAEKVMACFKmABCkTAEgSJkCQJAyBYAgR7DtpNJRaCmV5zqnTp2azSodZ1Rad9iwYdms0lxhaX5qy5Yt2az0saRUnjOF3lCaJe3o6Che29TUlM1WrVqVzWbNmlV5Y31IS0tLrbfQb3gyBYAgZQoAQcoUAIKUKQAEKVMACFKmABBkNGYnVRr5mD59ejbbvHlzNisd3ZZS9cehVRq5KV1bOmLJ6At9XelejBwRNmfOnGz20ksvVb1uT7n44ouz2f7771/1ujfccEPV1/ZHnkwBIEiZAkCQMgWAIGUKAEHKFACClCkABBmN6UWlE2cqncJSOjWmNP5S6dSY0rXDhw/PZqWTalJK6ZlnninmEHXaaacV81/+8pdVr33PPfdks/vvv7/qdXvKV77ylWz205/+tOp1b7vttmy2aNGiqtftjzyZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQZM60j+jq6qo676kj2FpbW6vKUjJnSs+76KKLinnpnql0VNo555xT1Z56yiGHHFLMr7/++mxW6f8tJaV5W97PkykABClTAAhSpgAQpEwBIEiZAkCQMgWAIKMxvWjKlCnZrHTEWkrlo9RmzJiRzTo6Oorrtre3Z7PSfisd7Qa7wl133ZXNWlpaiteWRkIWLFhQvHbbtm3FvCeMHDkym82fP7947QEHHFDVe06aNKmYL1u2rKp1d0eeTAEgSJkCQJAyBYAgZQoAQcoUAIKUKQAEGY3ZSWPGjCnms2fPzmalUZPS6S0ppbR27dqqskr7LZ3+UhotiJxEAd112GGHVX3to48+ms3mzZtX9boRpfGXe++9N5sNGzasuG7pfpw7d242+/3vf19cl+7zZAoAQcoUAIKUKQAEKVMACFKmABCkTAEgSJkCQJA5051U6ai0ao8ta2ws/7lm/Pjx5Y1ljB07tpiXPp7Sfh3Bxq5SumdK85WVvgZXrFiRzd58883itc3Nzdlsn332yWYXX3xxcd1p06YV85xKc+iXXnppNivNmbLreDIFgCBlCgBByhQAgpQpAAQpUwAIUqYAEGQ0Zhcr/Qh7afzlF7/4RU9sp6LSfqs99g12lchRf//617+y2cyZM4vXXnPNNdls4MCB2azSfkt56ci4n/3sZ8V177nnnmJOz/NkCgBByhQAgpQpAAQpUwAIUqYAEKRMASCooaubP3vulJDuaW9vz2al0zEq/TZMnz59l++n0vuWfs+nTp1aXLejo6O8sd1EZLSjJ/TF+7h0X9x+++3ZrNLH0lOf+9L7vvjii8VrFyxYkM3mzZuXzSqdckPP6s7XkidTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACCHMG2i5XmkUpZ6Si0lFJavHhxVddWmo9yBBu1Vjo+rDR7OWfOnJ7YTkoppeuvvz6bLV++PJtVui+2bdtW9Z7o2zyZAkCQMgWAIGUKAEHKFACClCkABClTAAhyBNsu1lNHsJU+/6VrKx2FdtVVV2Uz4y9xjmCD+ucINgDoBcoUAIKUKQAEKVMACFKmABCkTAEgSJkCQJA5011szJgx2WzVqlXZrNIRbKWZz6uvvjqb3XnnncV16VnmTKH+mTMFgF6gTAEgSJkCQJAyBYAgZQoAQcoUAIKMxkAPMhoD9c9oDAD0AmUKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAhq6Orq6qr1JgCgnnkyBYAgZQoAQcoUAIKUKQAEKVMACFKmABCkTAEgSJkCQJAyBYAgZQoAQcoUAIKUKQAEKVMACFKmABCkTOvQGWeckRoaGrL/bdmypdZbBAoeeeSRdN5556VRo0alQYMGpeHDh6epU6emDRs21HprVMl5pnVozZo16Yknnnjfa11dXWnWrFlpxIgR6a9//WuNdgZ0x+TJk9OqVavSlClT0ujRo9Nzzz2XrrnmmvT666+ntWvXpkMOOaTWW2QnKdN+4sEHH0zjx49Pl112WfrJT35S6+0ABatXr06HH3542nPPPd97bePGjenQQw9NkydPTrfccksNd0c1BtR6A+wat912W2poaEinnHJKrbcCVHDkkUd+4LWWlpY0atSo9Pe//70GOyLK90z7gbfeeiv95je/SUceeWQaMWJErbcDVKGrqys9//zz6ZOf/GStt0IVlGk/8Mc//jG99NJL6dRTT631VoAq3XrrrWnLli1p2rRptd4KVfA9037glFNOSR0dHWnr1q1pyJAhtd4OsJMef/zx1NramkaNGpVWrlyZ9thjj1pviZ2kTOvc66+/nvbbb7909NFHp2XLltV6O8BOeu6559K4cePSW2+9ldauXZuam5trvSWq4AeQ6txdd92Vtm/f7q94oQ69+uqraeLEiemVV15JK1euVKR1TJnWuVtvvTV95CMfSW1tbbXeCrATduzYkU444YS0YcOGtGLFinTwwQfXeksE+AGkOvbCCy+kFStWpEmTJqWBAwfWejtAN73zzjtp2rRpac2aNemOO+5IY8eOrfWWCPJkWsfa29vT22+/7a94oc58//vfT0uXLk0nnHBCevnllz/wjzScdtppNdoZ1fIDSHVs7Nix6cknn0zPPvusn/6DOnLUUUel+++/P5v733L9UaYAEOR7pgAQpEwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAjq9j8n2NDQ0JP7gH6pr/2bKO5j2HnduY89mQJAkDIFgCBlCgBByhQAgpQpAAQpUwAIUqYAEKRMASBImQJAkDIFgCBlCgBByhQAgpQpAAQpUwAIUqYAEKRMASBImQJAkDIFgCBlCgBByhQAggbUegMAlHV1dRXzzs7ObPbzn/88m1144YVV74n382QKAEHKFACClCkABClTAAhSpgAQpEwBIKihq9LPXL/7Cxsaenov0O908/bqNe7jntXU1FTMDz/88Gx20kknZbOzzjqruG7p62zjxo3Z7HOf+1xxXf6jO/exJ1MACFKmABCkTAEgSJkCQJAyBYAgZQoAQcoUAIIcwQawEwYPHpzNFi1aVLz2uOOO27Wb6Ybm5uZs1tbWVrx26dKlu3o7/ZYnUwAIUqYAEKRMASBImQJAkDIFgCBlCgBBjmCDHuQItv7nu9/9bja76qqreuQ9K/2+Vft19sADDxTzY445Zpe/Zz1yBBsA9AJlCgBByhQAgpQpAAQpUwAIUqYAENRvR2OOPfbYbFbp5IbSj4tPnDix6j2VNDU1FfNRo0Zls/vuuy+bPfHEE8V1b7rppmy2bdu24rVU1tfGB+rtPq431113XTE/66yzqlq3p0ZjKlm5cmU2O/HEE7PZq6++2gO7qR2jMQDQC5QpAAQpUwAIUqYAEKRMASBImQJAkDIFgKC6njP98pe/nM06Ojqy2T777NMT26no8ccfz2YPP/xw1esOHjw4m7W1tRWvffrpp7NZe3t7Nrvgggsq7gtzprubgQMHFvPS8W2l/2dNmDChuO63vvWtbHb44YcXry0pfb2MHDkym1Wab6835kwBoBcoUwAIUqYAEKRMASBImQJAkDIFgKABtd5AxKBBg7LZhz70oWy2devW4rotLS1V76nk7bffzmZvvvlm1es2Nub/TLTXXnsVr12wYEE2mzVrVjY79NBDi+t+9atfLebQH23fvr2YX3755VWtW2nU5Lnnnstmy5Ytq+o92TmeTAEgSJkCQJAyBYAgZQoAQcoUAIKUKQAEKVMACKrrOdO77747m82cOTObHXvsscV1K82K9TWdnZ3ZrNLHcvrpp2ezSZMmZbNFixYV1z3jjDOqvhag3ngyBYAgZQoAQcoUAIKUKQAEKVMACFKmABBU16MxJR0dHVVl/D9LlizJZgcffHDx2tLRbvfdd18227RpU6VtAf+fCRMmZLOGhoaq1y0d78j7+UwBQJAyBYAgZQoAQcoUAIKUKQAEKVMACGro6urq6tYvDPx4Nf1PU1NTMd+2bVs2a2lpyWb/+Mc/qt5TX9TN26vXuI/r04gRI4r5ypUrs9nQoUOrft+nn346m40bNy6bbd26ter37Iu6cx97MgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIKjfHsFGbZXmsg477LBs1t/mTGFXWLZsWTGPzJKWLFy4MJv1t1nSKE+mABCkTAEgSJkCQJAyBYAgZQoAQcoUAIKMxtDrWltbs9kdd9zRizuBvmPy5MnZ7OCDDy5eW+1Rfx0dHcV87ty5Va27O/JkCgBByhQAgpQpAAQpUwAIUqYAEKRMASDIaAy9bv369bXeAru5wYMHF/PTTz89mzU0NGSzSiMq5557bjYrjb80Npafezo7O7PZK6+8ks2uvfba4rp0nydTAAhSpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACCzJlSlX333bfqax988MFduBN2VyNGjCjmt956azYbOnRo8drhw4dns8icaUnp2tIcaaVrS7OkDzzwQHHd0ud427Zt2Wz79u3FdfsjT6YAEKRMASBImQJAkDIFgCBlCgBByhQAgozGUJXS2EFKKW3ZsiWbvfbaa7t6O/RTTU1N2ezMM88sXtva2rqrt1OXTjnllGx24403Fq9dsmRJNjvxxBOz2RNPPFFxX/2NJ1MACFKmABCkTAEgSJkCQJAyBYAgZQoAQUZjyGpra8tmo0ePLl576aWXZrPSaRPsfr797W9ns29+85vZrNLXIP9x0EEHZbPHH3+8eG3phBzez5MpAAQpUwAIUqYAEKRMASBImQJAkDIFgCBlCgBB5kx3c4MGDcpms2bNymavvvpqcd1rrrmm6j3Rv0ycOLGY19vXSmNj/hmks7OzeO3SpUuz2VNPPZXNzj///OK6ld63WqWPlffzmQKAIGUKAEHKFACClCkABClTAAhSpgAQZDSmnyuNvqSU0mWXXZbNjj/++Gw2Z86c4rovvvhieWPsNl544YViXjqSr6mpaVdvJ6w0hnLfffcVr503b142+/Of/5zNXnnlleK6F110UTGvVuljvfrqq7PZWWedVVy30tdEPfJkCgBByhQAgpQpAAQpUwAIUqYAEKRMASCooaurq6tbv7Choaf3Qg+4+eabi/mpp56azdauXZvNjj766OK6O3bsKG9sN9HN26vX9MX7uDSCdcstt2SzwYMHV/2eb775ZjEvneCycuXKbDZ79uziutu3by/mOQcddFAxX7JkSTYbPXp0Ve+ZUvnr5fXXX89mX/va14rrPvjgg1XvqRa6cx97MgWAIGUKAEHKFACClCkABClTAAhSpgAQpEwBIMicaT8wadKkbHbllVcWr33sscey2dlnn53NXn755Yr7wpxp1IQJE7JZZH6y0pFmpfnWvmjixInZbNmyZdmso6OjuG5pHrQ0i1t6z3pkzhQAeoEyBYAgZQoAQcoUAIKUKQAEKVMACDIaUycOOOCAbLZp06ZsVhp9SSmlM888M5utX7++0raowGgM1D+jMQDQC5QpAAQpUwAIUqYAEKRMASBImQJAkDIFgCBzpn3EF77whWK+fPnybNbZ2ZnNxo0bV1z3ySefLG+MEHOmUP/MmQJAL1CmABCkTAEgSJkCQJAyBYAgZQoAQQNqvYHdSekYtfb29uK1TU1N2ezQQw/NZkZfAHqeJ1MACFKmABCkTAEgSJkCQJAyBYAgZQoAQUZjetHJJ5+czT796U8Xr503b142+9vf/lb1ngCI82QKAEHKFACClCkABClTAAhSpgAQpEwBIEiZAkBQQ1dXV1e3fmFDQ0/vBfqdbt5evcZ9DDuvO/exJ1MACFKmABCkTAEgSJkCQJAyBYAgZQoAQd0ejQEA/jdPpgAQpEwBIEiZAkCQMgWAIGUKAEHKFACClCkABClTAAhSpgAQ9H+KJkIcgAY5/QAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 500x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize some data\n",
    "def visualize_data(data, labels):\n",
    "    fig, axs = plt.subplots(2,2, figsize=(5,5))\n",
    "    axs = axs.flatten()\n",
    "    for i in range(4):\n",
    "        axs[i].imshow(data[i].squeeze(), cmap='gray')\n",
    "        axs[i].set_title(labels[i].item())\n",
    "        axs[i].axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# visualize some data\n",
    "imgs, labels = next(iter(trainloader))\n",
    "visualize_data(imgs, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max:  tensor(1.)\n",
      "min:  tensor(-1.)\n"
     ]
    }
   ],
   "source": [
    "# max and min values of the data\n",
    "print('max: ', torch.max(imgs))\n",
    "print('min: ', torch.min(imgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "# device\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# image size\n",
    "img_size = imgs[0].shape[1]**2 # 28*28 = 784\n",
    "\n",
    "# hidden dimensions\n",
    "hidden_dims = 400\n",
    "\n",
    "# latent size\n",
    "latent_size = 20\n",
    "\n",
    "# output size   \n",
    "output_size = img_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([128, 784]) torch.Size([128, 1]) torch.Size([128, 20]) torch.Size([128, 20])\n"
     ]
    }
   ],
   "source": [
    "# define model\n",
    "model = VAE_GAN.VAE_GAN(img_size, latent_size, hidden_dims, output_size).to(device)\n",
    "\n",
    "# pass some data through the model\n",
    "imgs, labels = next(iter(trainloader)) #imgs.shape = [batch_size, 1, 28, 28], labels.shape = [128]\n",
    "\n",
    "# flatten images\n",
    "imgs = imgs.view(imgs.shape[0], -1).to(device) #imgs.shape = [128, 784]\n",
    "recon_imgs, y, mu, logvar = model(imgs) #recon_imgs.shape = [batch_size, img_size], mu.shape = [batch_size, latent_size], logvar.shape = [batch_size, latent_size]\n",
    "\n",
    "print(recon_imgs.shape, y.shape, mu.shape, logvar.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 1/100 [00:06<10:35,  6.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 10459.0830\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 11/100 [01:08<09:12,  6.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10, Loss: 6784.9341\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|██        | 21/100 [02:10<08:10,  6.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 20, Loss: 6043.5181\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 31/100 [03:13<07:10,  6.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 30, Loss: 5740.3398\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 41%|████      | 41/100 [04:14<06:06,  6.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 40, Loss: 6151.8838\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 51%|█████     | 51/100 [05:17<05:04,  6.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 50, Loss: 5672.2002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|██████    | 61/100 [06:19<03:59,  6.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 60, Loss: 5459.8911\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 71%|███████   | 71/100 [07:19<03:00,  6.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 70, Loss: 5920.9580\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 81%|████████  | 81/100 [08:22<01:58,  6.23s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 80, Loss: 5903.0029\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 91%|█████████ | 91/100 [09:24<00:55,  6.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 90, Loss: 5808.8423\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [10:20<00:00,  6.20s/it]\n"
     ]
    }
   ],
   "source": [
    "# training loop\n",
    "epochs = 100\n",
    "for epoch in trange(epochs):\n",
    "    for i, (imgs, labels) in enumerate(trainloader):\n",
    "        # flatten imagessamples =\n",
    "        imgs = imgs.view(imgs.shape[0], -1).to(device)\n",
    "\n",
    "        # forward pass\n",
    "        recon_imgs, y, mu, logvar = model(imgs)\n",
    "\n",
    "        # loss\n",
    "        loss = model.loss_function(imgs, recon_imgs, y, mu, logvar)\n",
    "\n",
    "        # backward pass\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        # update weights\n",
    "        optimizer.step()\n",
    "\n",
    "    # print loss\n",
    "    if epoch % 10 == 0:\n",
    "        print('Epoch: {}, Loss: {:.4f}'.format(epoch, loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "decoder_out.shape: torch.Size([64, 784])\n"
     ]
    }
   ],
   "source": [
    "# generate some data\n",
    "new_samples = model.sample(num_samples=64)\n",
    "print('decoder_out.shape:', new_samples.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAekAAAHqCAYAAAAgWrY5AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAuvElEQVR4nO3deXDV5fn+8SdCAoSdsAQIEBDZCZuKgKAUiwtC1Y6jIrVYtThddKplOjpWW9e2VqettLZjFbcBFBXcZaRGkE0BAQHZlJ2wBcK+L78/fvOdjt7XYz8hgdzn5P3685pznnNCPufcfsYrz5Nx8uTJkwEAALhzVkW/AQAAoDGkAQBwiiENAIBTDGkAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMAphjQAAE5VTfrAjIyM0/k+UImc6iZ3XIMoL2XZaJHrEOUlyXXInTQAAE4xpAEAcIohDQCAUwxpAACcYkgDAOAUQxoAAKcY0gAAOMWQBgDAKYY0AABOMaQBAHCKIQ0AgFMMaQAAnGJIAwDgFEMaAACnGNIAADjFkAYAwCmGNAAATjGkAQBwiiENAIBTDGkAAJyqWtFvoKJkZGTIvFq1aiZr2LChyfLy8kyWk5Mj16xRo4bJlixZYrIdO3aYrKSkRK557NgxmQNAaVStqsdAz549TTZq1CiTZWZmmmzZsmVyzfHjx5tsw4YNJjt+/Lh8fmXEnTQAAE4xpAEAcIohDQCAUwxpAACcqhTFMVXc+uEPfygfe9ddd5msTZs2JqtevbrJzjor+X/znDx50mRHjhwxWVFRkXy+KnB88sknJqOAAeD/qO+o7t27y8f+61//Mlnr1q1Npkq4F154oVyzZcuWJrvvvvtMpkq0lRV30gAAOMWQBgDAKYY0AABOMaQBAHAq46RqMKkHRnbo8kYVI4YOHWqyMWPGyOc3b97cZOpnT/jPFpX03zP2Olu3bjWZKsPNnj078Zpnyqm+fqpcg2Wlfk6VValSJVEWgv5cHDp0yGQnTpxI8hZTXlk+A6l8HdauXdtkjz76qHzsjTfemGjN3bt3m+zgwYPyseqae+qpp0z24osvmiwdS7BJrkPupAEAcIohDQCAUwxpAACcYkgDAOAUQxoAAKfSbltQdTbqkCFDTNaoUSP5/KRNbtWCPXr0qFxTnf2szq1WzdxYk1SdcT1s2DCTzZ07N/H7xOmjft+x88c7depksm7dupmsVatWJotd1+p3/uGHH5rs7bffNplq74ZQ8X8lgO+mvguvvPJKkw0fPlw+X7Wply5dajL1FyR9+/aVa6otlq+77jqTjRs3LtH7qQy4kwYAwCmGNAAATjGkAQBwiiENAIBTaVccU0Wr9evXmyxWnlLlLVX82rBhg8kKCwvlmlu2bDFZly5dTNavXz+Tlabgpsoa6ixtimOnV5MmTUx20003mezqq6+Wzz/77LNNlp2dbbKsrCyTxbYFVfkVV1xhspo1a5ps7Nixcs3Dhw/LHD6okpbaAlRdWyGEsHz5cpM98cQTJlPfhXl5eXLNnj17mqxjx44mi13HlRF30gAAOMWQBgDAKYY0AABOMaQBAHAq7YpjR44cMdmECRNM1qdPH/l8VXjYvHmzyZ588kmTrV69OslbDCGE0LVrV5OVZgcppX79+iZT5wij/KjzeUeOHGmyO+64w2SqpBVCCDt27EiUZWZmmqxOnTpyzQYNGiTKVJls0qRJck11pjkqhiqSfu973zNZ48aNTRY7+1l9b06bNi3R+1Fl2RD0NVuvXj2Tqe/hlStXJnrtdMM3OAAATjGkAQBwiiENAIBTDGkAAJxKu+KYOj5v7dq1Jrv77rvl89VuUeqoPlWMUKW1EHSZJzc312QtWrQwWeyoSvVzHjp0yGQUx8qHKryEEMKAAQNMdtddd5msVq1aJluyZIlc89lnnzVZUVGRydSuTOoaCiGEm2++2WRt27Y1mSofqmM24Yv6nujevbvJ1DWzf/9+ueb8+fNNduDAAZNVr17dZLFdzNTnSH1vxgqQlRHf4AAAOMWQBgDAKYY0AABOMaQBAHCKIQ0AgFNp1+5W1HnQqvEdQgj79u0zmWoqqta02iIyBL0FqWoAq20aY1QjcubMmSaLbfmH0on9bm6//fZEjy0pKTHZX/7yF7nm1KlTTabOblbX5bp16+Sa6tq85pprTKauf9XeDUF/Bk6cOCEfizNPNbHVWfKx36/arlNtZav+IqCgoECuqa4P9Z7UX69UVtxJAwDgFEMaAACnGNIAADjFkAYAwKlKURxTVJksBF1YyM/PN1ndunVN1qNHD7nmqFGjTKZKGaUp4nz99dcme+mll0ymfs7SbDVaGal/n9atW8vHnn/++YnW/Oijj0w2e/Zs+Vi1Da3azlEVblTxK4QQFi5caDK1bWSzZs1MprY+DUGfs7537175WJxe6nti3LhxJlNlwZycHLnmtddeazL12VDby6rtlUPQ3zHqmlXnXse25lWfg3TCnTQAAE4xpAEAcIohDQCAUwxpAACcqhTFsVhRSlElittuu81k5557rsliZQm1o496T0lLFSHoM4dXrlyZaE0KYt9NFfjy8vLkY1Wha+fOnSZ74YUXTLZjx47E7+n48eMmU+cAxwqRVavaj/rGjRtN1rNnT5MNHTpUrqnKaCqLvSecXosWLTJZYWGhya666ir5/IsuushkqtClvssaNmyY4B3+f+rzduGFF5pM/TwhhFBcXGyydNrFjDtpAACcYkgDAOAUQxoAAKcY0gAAOFUpimOqmBDbZedPf/qTydSOS9WqVUv0OqWhCjaxXakmT55sMnV8JaWd8lGrVi2Zq5KYKtKo58d2UEpadFS7TMV2qFPvc/r06SarUaOGyTp27CjXHDRokMlUGW3Lli3y+Ti9VHlq9OjRJqtTp458vjpit1u3biZTpUR1jGoI+tpWuzd27tzZZLEymtqhL512IeNOGgAApxjSAAA4xZAGAMAphjQAAE4xpAEAcCrt2t2qPdigQQOTPf744/L5F198sclUk7s01HZ0qoWrtnlcv369XFNtF5pOjUZvlixZInO1VWGbNm1MNnDgQJPt2rVLrrlhwwaTqd+3ui7V9qGx1/rkk09MtmfPHpPdeeedck3VwM3NzTXZ1q1bTZaqWzSmOrUV7YMPPigf++tf/9pkajvkli1bmkxtlxuj/ipFfT/GGuOqXZ5OuJMGAMAphjQAAE4xpAEAcIohDQCAU2n3f9zVtoZXX321ydSWhiHoEkLSM5ljpZ3Dhw8ner4qUKjzW0MIoV69eiZTWz+i9FRp5euvv5aPfe6550z2s5/9zGRqa83Y9TJv3jyTrV271mRqO0S1LWcIya9B9XMuXrxYrqm2iFQlouXLl5vs0KFDck2cXup3HitFPvHEEya75557TFa/fn2TxYpj6vUPHjxoMrVVqTpjOgRddty7d2+i104F3EkDAOAUQxoAAKcY0gAAOMWQBgDAqZQujqmSlyquDBkyxGSxM1TV+ctqJ69NmzaZTJV7QgihevXqJlNFIlV6U+WcEELIz883WazchNJRBRNV0gpBn8msdo7r1auXyVq3bi3XVOeSq+LXmjVrTKbKh9+Vf1tpyofqs6Z2IVP/RhTH/IidQa52uUu60+HKlSvlmur6UjuJtW3b1mQ//vGP5Zrq87Zq1SqTxX5O77iTBgDAKYY0AABOMaQBAHCKIQ0AgFMpXRyrXbu2yW699VaT9ejRw2Sx3Z5U+eqvf/2ryZYtW2YytctNCCEUFBSY7LHHHjOZ2kWsVq1acs2cnByZ4/SIlU5UaWXGjBkmmzt3rsliv0N1BKU6YlDt1BS7rpOWZlRpLnYEqjoWtkuXLibLzMxM9NqoGLGduLKyskzWokWLRM9XR7iGoAtl/fv3N5k67lUVKkvznlIVd9IAADjFkAYAwCmGNAAATjGkAQBwiiENAIBTKd3uVtsqqrOj69ata7J169bJNe+++26TLVy40GRqS9LYGarqTGjV4FVtWbWlaAghNGrUSOaoeKpJrZrYamvZEOIt1iSvU9ZWq2qHx5rhzZo1M5m6htW1Dv/UtaD+okZtORs717xhw4Yma9CggcnUd2ns2t66dWvix6Yi7qQBAHCKIQ0AgFMMaQAAnGJIAwDgVEoUx2JFmt69e5tMlbRUyWvXrl1yTbUFoiqotWrVymSqSBNCCCNGjDBZ/fr1TaZ+TlU4CkFvS4rUEiu3xLb2PBNU8UtdqyHoUqP6mdSa8E99R6qSWNLtYUPQJTGVHTt2zGTr16+Xa06bNs1kFMcAAMBpx5AGAMAphjQAAE4xpAEAcColimOxnbzUrlvq7Fr1fHXGcwghPPLIIyZT5zyrskNsZyV1JrQqs6lSxvz58+WaK1asMFk6lSVQMdTuYnv27JGPVfnatWtNFjuPGr6p76M1a9aYTO3o2K5dO7mmeqz6Lty9e7fJFixYINdkxzEAAFAhGNIAADjFkAYAwCmGNAAATqVEcSx2VJ4qpKhMlcnUkWshhNC3b1+TJd0xKVZWULnaSWzVqlUm+/3vfy/X3LJlS6L3BJSG2u1s6dKl8rHqet2/f7/JYp9f+LZv3z6TTZ061WRt2rQxWX5+vlwzOzvbZAcOHDCZuubGjRsn1ywuLjYZxTEAAHDaMaQBAHCKIQ0AgFMMaQAAnGJIAwDgVEq0u2Pn606cONFkHTt2NNmQIUNMVqdOHblm7Ozqb1PtQbWNXgi6ffjJJ5+Y7OGHHzaZ2v4zhIo9cxjpS11XCxculI8tKioyWUlJiclUuzv2FxPp1MpNdepamDFjhsn69+9vsry8PLmm+quWjz/+2GT//Oc/TRbbIlmdPZ1OuJMGAMAphjQAAE4xpAEAcIohDQCAUxknEzY1km6NWdFUIWzgwIEm69Onj3x+586dTaa2EFXFrylTpsg1v/zyS5Pt2rXLZJVl+8RTLQelyjWYbnJzc2U+YsQIk/Xs2dNkDz74oMlihcgzVRwry+twHX5TlSpVTBYr5qptQXfu3GmyQ4cOmSwdS4VJfibupAEAcIohDQCAUwxpAACcYkgDAOBU2hXH4B/FMb/Uv3GLFi3kY3v37m2yli1bmmzdunUme+utt+SasV37yhvFMXhAcQwAgBTGkAYAwCmGNAAATjGkAQBwKiWOqgRwZqgiS2ZmpnysOoJ18ODBJuvSpYvJ3nvvPbnmmSqOAamCO2kAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMAptgXFGce2oOlL/Y5UVtFnp7MtKDxgW1AAAFIYQxoAAKcY0gAAOMWQBgDAqcTFMQAAcGZxJw0AgFMMaQAAnGJIAwDgFEMaAACnGNIAADjFkAYAwCmGNAAATjGkAQBwiiENAIBTDGkAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMCpqkkfmJGRcTrfByqRUz0dlWsQ5aUsJ/RyHaK8JLkOuZMGAMAphjQAAE4xpAEAcIohDQCAUwxpAACcYkgDAOAUQxoAAKcY0gAAOJV4MxMAZ4baLOOss+x/T9eoUcNkhw8flmsePXq07G8MwBnHnTQAAE4xpAEAcIohDQCAUwxpAACcYkgDAOAU7W6ggsSOPMzMzDRZz549TdarVy+Tvfnmm3LNoqIik504ceJ/vUUAFYw7aQAAnGJIAwDgFEMaAACnGNIAADhFcQw4A1RJrEmTJvKxI0eONNmwYcNMduDAAZPt3btXrjl58uREjz158qR8PoCKwZ00AABOMaQBAHCKIQ0AgFMMaQAAnKI4BpwBTZs2NdnTTz8tHzt48GCTVa1qP6rqjGhVJgshhDlz5iR67LFjx+TzkR7UueSleezx48dNRtnw9OJOGgAApxjSAAA4xZAGAMAphjQAAE5RHHNE7UpVpUoV+dhq1aqZTB09qMpFlINOr+rVq5vsqaeeMtnll18un6+OqlTlnIMHD5osNzdXrtmvXz+THTp0yGSbN282mbqGUDGysrJkrq652rVrmyw/P99kderUkWs2btw40evv2rXLZOraDEEXz1avXm2yLVu2JF5TXZ/pVGbjThoAAKcY0gAAOMWQBgDAKYY0AABOMaQBAHCqQtvdqrms2n8xqg2t2oeq6Rdr/6mG9OloCqptHmvWrGmy3r17y+d3797dZIWFhSabO3du6d8cymTAgAEmU01u1eKO2b9/v8kWLlxostdff10+f8WKFSZTTXB1rW/atEmuqT4rKD+NGjUy2X333Scfq9r7OTk5JlON7djv8fDhw4keq77HYy30I0eOmGz37t0mW7NmjclmzZol13zhhRdMphrjqXq9cicNAIBTDGkAAJxiSAMA4BRDGgAApyq0OJa0JBbbGrNWrVom69q1q8nUNph79+6Va6pcPV8VI2LbbaqCW8uWLU122223mezKK6+Ua65bt85kH3zwgcnKWnpTBbe8vDyTrV27tkyvk6rUdaDKPWob1xi1XeeECRNMNmbMGJOp7RRD0IUd9flRnzX1uBBC2Ldvn8lStZxT0WrUqGGyO+64w2SxrWRbtWplsqTFRPX9VBrqmol976i8QYMGJlOludatW8s1V65cabKNGzeaTH2uUgF30gAAOMWQBgDAKYY0AABOMaQBAHDK3XnSqsSQnZ0tH3vJJZeY7J577jGZKnSpYkEIycsw7du3N1msCKd2z+nQoYPJevToYbJYaW79+vWJXqesxTH1b1dZS2KKKgA2a9Ys0XNj18tnn31msmeffdZk6vd94MABuWbS86hVSSz2+VM7UqmCWjqd7Xu6FBQUmEwVoqZNmyafr4pjqihVVFRkstjvp6SkxGRqd7A9e/aYLFZaO++880zWv39/k6nimHrtEPR3eazEm4q4kwYAwCmGNAAATjGkAQBwiiENAIBT7opjaoer5s2by8eOGjXKZOecc47Jdu3aZTJVcAlB7zimShlqh6BYsUEdQanKRapscfToUbnm448/brLYLmoou9iuTPn5+SarV6+eyVT5cNu2bXLNp59+2mSrVq0yWdKjBEMI4ayz7H+Pq8+a2hktVl6sX7++yVQZThWQYqW5pCUz9Z5iP3sqFNeWLl1qsq+++spkU6ZMkc9XhUFVnirNUcDqmlGZumbUkZgh6N0K1feW+v2+9tprcs0FCxaYrDQ/p3fcSQMA4BRDGgAApxjSAAA4xZAGAMAphjQAAE6lRLs7doaq2lpTWbZsmcmee+45+Vi1JaOi2ov79++Xj61du7bJJk6caDLVQt20aZNc85133jEZZ/mePqrVGkIIubm5JlNtV9U2Xbx4sVxz3rx5Jku63ab6/ISgr0H1lxBqe0rVIo+9J7Wt6PTp0022YcMGuaa6htW/nWrQFxcXyzVTgdqOOPZ9klRZW+3qmlfnp6vfRb9+/eSaV111lclycnJM9vrrr5vs3//+t1xTtcNTodGfFHfSAAA4xZAGAMAphjQAAE4xpAEAcMpdcUwVE/r06SMfq7bb3L59u8nUNouxc1mTlhDUec4xqmRWvXp1k6mCzIQJE+SaZS2VoHRi24Kq83EVdebu/Pnz5WOTbu+qtqZt0qSJfOyAAQMSZernVOcah6A/f4pa8/PPP5ePVeWvrVu3JnqddCoLhVDxP496fVUSu/baa002fPhwuaYqFs6cOdNkDz30kMnUd3sI6V+Y5U4aAACnGNIAADjFkAYAwCmGNAAATrkrjmVnZ5ssdjao2glJ7dC1du1ak6nzV7/rtcpCFWySFsfeffdduWZFl0oqG3XWdwi6SKN+j6r8tGTJksSvr17n4osvNtkNN9wgn9+2bVuT7dixw2SzZ8822caNG+WadevWNVnnzp1NpnZBixUvi4qKTKY+K7EzrlF+1I5j6jzovn37mkxdWyGEMGPGDJO9/PLLJtu8ebPJ0r0gFsOdNAAATjGkAQBwiiENAIBTDGkAAJxyVxw7duyYyWIlL1UcUzuWqSMtY2UY9VrqdVRxK3acoSpWKDt37jSZOsIuBP1zqp2dVNlC/RuHkLyMVhlLa6q8FIIuZKl/n127dpkstrOYKuf85je/Mdmll15qsljBTb2WOipz+fLlJouVvNq0aWMydXyl2kVs27Ztcs3Yv/O3lZSUJHoc/rfYbnqNGjUy2bBhw0zWvHlzk8V+P+rIVJWpsm3s85LuhTLupAEAcIohDQCAUwxpAACcYkgDAOAUQxoAAKfctbtVK/DVV1+Vj23QoIHJOnbsaLK7777bZDfddJNcc+7cuSZbsWKFydTZpmrrxhBCuOuuu0ymWriq2XrvvffKNT/88EOTHTx40GRr1qwxmfp5Ys8/evRooizdG9+tW7eWuWq2qi0rVaaa4SHo1vagQYNMps7mjTX3VZtabdG4YcOGxGuqv0ZQr/Pll1+aLNbIVZ9p9frqLy7U9YtvUn+BEvve+sEPfmCyK6+80mTqXHPVDA8hhKZNm5pMfZfm5OSYTH3nhaCvw3TCnTQAAE4xpAEAcIohDQCAUwxpAACcclccUyWRjz/+WD5WnR192223mWzAgAEmixWB2rVrZzJV6FLbH8bOom7ZsqXJ1LaeqgzTvn17uaYqIn300UcmU+e6quJXCHp7QFUIS/eSmKK2LgxBXweqFKXOVI5tc6g+A+p6Uddb7BxfdW2oLUCVVq1ayVyVgHr06GEydb2obR9D0Nvgqn8nVTaKbcub7ttGxqh/D/Xv3qlTJ/l8dV65uubWrVtnsth1qEpq/fv3N9m5555rstjvd/LkySZLpxIhd9IAADjFkAYAwCmGNAAATjGkAQBwyl1xTFGFqhD0TkajR482WbNmzUymdiYLIYSGDRuaTO2ClJ+fb7KCggK5pireqDKL2gls7Nixcs1Vq1aZbP78+SZTpZtYwU2pjCWx0kh61rgq68XOflZnmqsijLqG1LnVIejCzvXXX28yVSxq0qSJXLNx48YmU6W3unXrmixWxFO7qKmz31VWWQtiIehSVXZ2tslUMXbo0KFyzWrVqpnsrbfeMllhYaHJ9uzZI9dU37u33HKLybp06WKykSNHyjVVgXjmzJkmK833nifcSQMA4BRDGgAApxjSAAA4xZAGAMCplCiOxaiiiNoBSu2Is379ermm2nVLlTLUjl8XXXSRXFMVylRp6O9//7vJJk2aJNdU5aLYkYIou+LiYpmrYp4qjqnduYYPHy7XVKUqteudulZVmSuEEK644opEz1fHDsZ2elLPVzuBqX+72PGIW7ZsMVlRUZHJVBkN36RKeGr3w1gx9+233zbZu+++azJVVlTfjyHo7yh1PLH6fmzRooVcUx2LqUpv6jszFXAnDQCAUwxpAACcYkgDAOAUQxoAAKcY0gAAOJXS7e6kynomstpOTrUU1VaFIeimo9om8tNPPzWZOl83hMq9BWJFWL16tczVVq6dO3c2WV5ensmaN28u11StaXUNqceprSBj1DWkPheHDh2Sz9+8ebPJ1L+Huq5nzZol11QN3P3795ss9tcZlZVqM6vtiFXje9GiRXLN6dOnm2znzp0mU9eMujZD0FvMqr98UOdRv/baa3JNtR1yrLGeiriTBgDAKYY0AABOMaQBAHCKIQ0AgFOVojh2OqiyROyMXLXFnTpvVZUyKIj5ECtPqfO+e/fubbK2bduarGpV/fGLlW6+TZUXY1vDqnz37t0mU2Ww9957T645efJkk61Zs8Zkaqve2Nm+Sc/8ZVvQ/01976htWxcuXCifr64P9b2nto09++yz5ZoPPPCAyVq3bm0yVTacMGGCXFOdJ52qZ0cr3EkDAOAUQxoAAKcY0gAAOMWQBgDAKYpj5ahmzZoyV+UvlaldyOBDbIc6VXC5+eabTTZ69GiTderUSa6pdo9S5at58+aZLFYCKiwsNNnWrVtNps4GPnjwoFyzLKXGWDku6e5VfFa+Ken3idpxLDc3V66pdnrLysoyWffu3U32xz/+Ua7Zq1cvk6kymzq3WpUSQ0j/a4E7aQAAnGJIAwDgFEMaAACnGNIAADhFcawcqVJFjDpKTe3cA9/U73Hu3Lkmu+GGG0xWvXp1uaa6DlQ5Ru2CFiu4leZo1jMh9n6Sfgb27t1bnm8n5anrY8uWLSYrKCgw2dChQ+WaxcXFJlM7511++eUma9eunVxT7WL2zDPPmGzixIkmixUY0x1TAQAApxjSAAA4xZAGAMAphjQAAE4xpAEAcIp2dzlSWzeGoM+OVmegqvOo1RnTIfhr6+K/2Ab21HFdnxp1faktXtVfI1x22WVyzZYtW5pMbSuqfPHFFzJX24W++eabJoud314ZcScNAIBTDGkAAJxiSAMA4BRDGgAApyiOlSO1jV4IIezbt89kmZmZJsvLyzNZUVGRXDNWUgNSGcWx8qO2WC0pKTFZbItVVTJTZ0xPmjTJZK+88opcc+nSpSajQPnduJMGAMAphjQAAE4xpAEAcIohDQCAUxTHytG2bdtkvnXrVpOps16rVrW/jgYNGsg11fMB4P+oXci2b99usueee04+v0WLFiYrLCw02erVq02mdlmMvSd8N+6kAQBwiiENAIBTDGkAAJxiSAMA4BTFsXKUkZEh8xUrVphs8uTJJps/f77J1K4/APC/7N6922SfffaZyebMmXMm3g5OEXfSAAA4xZAGAMAphjQAAE4xpAEAcIohDQCAUxknEx7gGmsu47+ysrJkrs5+rlatmslUk1ttFRpCCMeOHSvlu/PjVM8M5hpEeSnLudVchygvSa5D7qQBAHCKIQ0AgFMMaQAAnGJIAwDgVOLiGAAAOLO4kwYAwCmGNAAATjGkAQBwiiENAIBTDGkAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMAphjQAAE4xpAEAcIohDQCAUwxpAACcYkgDAOBU1aQPzMjIOJ3vA5XIqZ6OyjWI8lKWE3q5DlFeklyH3EkDAOAUQxoAAKcY0gAAOMWQBgDAKYY0AABOMaQBAHCKIQ0AgFMMaQAAnGJIAwDgFEMaAACnGNIAADjFkAYAwCmGNAAATjGkAQBwiiENAIBTDGkAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMCpqhX9BoB0k5GRYbIqVaqYLCsrSz7/6NGjJjt27JjJTp48eQrvDvgmdb3GcnXNcR2eXtxJAwDgFEMaAACnGNIAADjFkAYAwCmKY0ACZ52l/3u2du3aJmvYsKHJ+vbta7J+/frJNQ8fPmyyZcuWmWzx4sUmW7BggVzzwIEDMkd6UCWv7OxskzVt2tRkPXv2lGv26NHDZM2aNTPZwYMHTbZu3Tq55sqVK002c+ZMk23fvt1kx48fl2umO+6kAQBwiiENAIBTDGkAAJxiSAMA4FTGyYTbxcR2pcHplZmZKfNzzjnHZH369DHZokWLTDZ//ny55pnaOehUX6cir8HYazdv3txk999/v8l69+5tsiZNmsg1q1a1fU51HezZs8dkGzdulGvOmzfPZH/4wx9MtnnzZpOdOHFCrpnKynKte/wuVLvXtW7d2mS33HKLydT3Rgj6OlQ73yV9P7E158yZY7J33nnHZB9//LFcUxXXUkWS65A7aQAAnGJIAwDgFEMaAACnGNIAADjFkAYAwKlK2+6O/Txq+8ek7dbYP6VqNLZv395kv/vd70w2aNAguWa1atVMtnPnTpP99re/NdkLL7wg16TdHad+hyGEcOmll5rsV7/6lcm6du1qMnXGdAj6GqxevXqi58f+jdSWimrrxdGjR5vszTfflGseOnRI5qkgVdvdse1pCwoKTKb+yqBly5Ym++qrr+SaartOta3nrl27TNamTRu55sCBA02mPhu1atUy2ZIlS+SaDz/8sMlWrFhhsqTN9DOJdjcAACmMIQ0AgFMMaQAAnGJIAwDgVNqdJ622o2vRooXJ1LaaIeiCzu7du03WoUMHk/Xv31+uqc5g7dSpk8nq169vslhR5MiRIyZTZwZPmTJFPh/lQ5WnVq1aZTJVWlG/7xD0NodqW9A6deqYTBWDQgihRo0aJmvcuLHJVAmnS5cucs2//e1vJisuLjbZmSokVgZ169aV+c9//nOTnX322SabPn26ycaPHy/XVGeTq7POFbUdcQi60DVkyBCTDRs2zGSxEq265tX2p6og57FM9m3cSQMA4BRDGgAApxjSAAA4xZAGAMCplC6OZWdnm+yll14ymSppxXZ7+vLLL01WWFhoMlXKUAWxEHTZQp2hesEFF5hMFX5C0Lv8qCKPehxFntJTO3aFoAsy69atM5k62zd2Du6WLVtMdvToUZOpkuP5558v17zmmmtM1r17d5M1bNjQZKNGjZJrqsKO2rFs27Zt8vn4bmpns1gxMDc312RqB8I33njDZOqs8RB0OTUpdb2GoMtbkyZNMpkqzI4YMUKu2a5dO5OpIt1DDz1kMrXrXgi+viO5kwYAwCmGNAAATjGkAQBwiiENAIBTKVEcUwWZEEJ4//33TdavXz+TqV1lFi9eLNccM2ZMoseqUkfs6D5VQujWrZvJVOktVhwrKSkx2bRp00yWdIcgfLdYkWTHjh0m27Nnj8m2bt1qstj1oo5GTVpkWbt2rczV7lGPPfaYyS666CKTxa5BVVJTJR51rcaKeOrnVJ+1pMfHppsGDRrIXB2lqq4vtXti7HdRFrHfj9oVcf369SabOHGiydq2bSvXvOKKK0x24YUXmqxp06YmU/8eIfj63uROGgAApxjSAAA4xZAGAMAphjQAAE4xpAEAcMpdu1s1Of/xj3/Ix6omt3q+OsN08ODBcs2ybKOpXjsEvcXdxRdfbDK1rWhsTbV9ozprVm1/Gvt5PG2FlyrUv5lqy6otQMvaqlXXRuz8cdWMbd++vcnUX1LErsG9e/eaTF2DqpGsnhuC/kuM09E+TlX79++XeU5OjsnU7+KSSy4xmfp+/K7X+jZ1fcS2Xc7KykqUqe1P1fnnIejtS1Wmtk5dvny5XFP9TBX1/cidNAAATjGkAQBwiiENAIBTDGkAAJyq0OKYKrl07NjRZDfeeGPi56tCijqnOXaWb1KlKUuo86zvvPNOk6mt/WJlBVXwufXWW02mijhqi8gQdEGH0k75UNdGZmamfGzSQli1atVMNmDAALnmI488YrI2bdokep9qm9MQQpgxY4bJNmzYYDJ1Dca2XSzLlqjpRv3canvZEELYt2+fyVQx8Ec/+pHJvvjiC7nmzJkzTZb0e7NevXoyV9+FvXr1MpnanjYvL0+uqc6uVtt9qi1zY2dme7rmuJMGAMAphjQAAE4xpAEAcIohDQCAUxVaHFPFmUcffTTR40LQ/3P/9ddfN1lpzgZVxZlWrVqZLOluTSGE8Itf/MJk9evXT/R+Yrs91axZ02Q33HCDyS699FKTXX/99XLNjRs3mmzbtm0mUyUV/Jf6nTVv3txkqrgVgi7IqLNwe/bsabLzzjsv8Zrqfarf7YsvvijX/POf/2yy4uJik6kyGIXEU1NUVCTzDz74wGSdO3c2mdr97fvf/75cc/PmzSZT3wdNmjQx2U9/+lO55qBBg0ymdkZT1LnTIegy3ezZs02mSsWeCmIx3EkDAOAUQxoAAKcY0gAAOMWQBgDAqQotjqmdYvr06ZP4+WonI1XGKSgoMFmdOnXkmk8++aTJVElM7QB16NAhuWbt2rVNFiuEJaUKbmr3HLVDUKyAoUohZd2ZrTJSO8f17t3bZMOHD5fPV49VRUF1DcSOqlS/xzVr1pjsqaeeMtnLL78s11TXeyoUcVKZ+s4MIYRnnnnGZOr6UDs6Tp06Va5Zq1Ytk6kjJEeMGGEyVVgNQR9LeeDAAZN99tlnJosdK6nekyobxz4b3qXmuwYAoBJgSAMA4BRDGgAApxjSAAA4xZAGAMCpCm13qyboRx99ZLLLLrtMPl+dXavO2L3//vtN1qFDB7mm2trzq6++Mtm6detM1qhRI7lmaRrr3xY7y3f8+PEme/jhh022fft2k6XCGarppm3btibr3r27fKzaulH9NYDKYm181ehX23q+8847Jov91QLXix8lJSUmU039du3aJV5Tta7VOdHqr19UMzwE/X00duxYk3344YeJXjuEEH7yk5+YrG/fviZ74403TLZp0ya5pqdrmztpAACcYkgDAOAUQxoAAKcY0gAAOOWuOHbTTTeZTG2rGUII+/fvN1l2drbJVFEqVp5S2+4lLRGoclAIIaxatSrR81VpLnbWqzqjFz6o7WrVNfD111/L59eoUcNkqhAZO2ddUVsixraYRHpQJUJV3OrVq5d8vioMqseqAuO8efPkmg888IDJ1NnPSc9kD0GfE921a1eTqff++eefyzUpjgEAgP+JIQ0AgFMMaQAAnGJIAwDgVIUWxxRV6NqxY0fi58d2RzoT3n///cSPPX78uMkGDx5sMgpiqUf9bgsLC02Wl5cnn7927VqTqUKk2rGsadOmcs2cnByTderUyWTTp083WWzXO1WQgx+q/FSnTh2T3XjjjYnXzM/PN5kqRY4ZM0Y+f8aMGSZTnxdVHCsuLpZrqoKcKkq2atXKZOrs9xDixeKKwJ00AABOMaQBAHCKIQ0AgFMMaQAAnHJXHEsVqpjQpk2bxM/fvHmzyVSBAqlHFXYOHz5sslmzZsnnL1261GQ1a9Y0mSpuDRw4UK6pHquOVlW7++3cuVOuqUqNFB39UL+Lzp07m+yCCy6Qz1fHTaojep955hmTzZkzR66ZtGyoPkPq6MwQQlixYoXJunXrZrJYAdI77qQBAHCKIQ0AgFMMaQAAnGJIAwDgFEMaAACnaHefonPPPddkqvEdc91115Xn24Ej6jpQrWl1bnQIektC1WxVLXB1jm4I+jxqdR579erVTUZjO32o7WWbNWsmH6sa1suWLTPZp59+arLTsa1m7DrMysoymXrvu3fvNlkqbG3LnTQAAE4xpAEAcIohDQCAUwxpAACcojh2ip5//vnEj1WlH1W2QHpQxbEqVaqYrHHjxvL5qmRWt25dk7Vv395kmZmZck31+mqrUVUwi5VrVDkHvm3YsCHxY9U1ozJ1HZT12ijNtssDBgwwWWmuY++4kwYAwCmGNAAATjGkAQBwiiENAIBTFMcSUCWGVq1aJX7+559/bjLOjq5c1Nm8l1xyiXxshw4dTKbOflZlsqpV9UdanQldUlJisvXr15ssttMTxbHUo855XrRokXys2r1OndPcsWNHky1YsECuefToUZNlZGSYrEWLFia788475Zr5+fkmU2dMv//++yZLhWuYO2kAAJxiSAMA4BRDGgAApxjSAAA4RXEsgdzcXJOpMlmsYDNr1qxyf09ILar41a5dO/nYgoICk6nj+NQOSnv37pVrqmtw/PjxJlMFs1Qo1yAZdTzpyJEj5WMnT55ssiZNmpjsscceM9kbb7wh19y8ebPJ1M55Q4cONVnDhg3lmrt27TLZlClTTLZ9+3aTpcK1zZ00AABOMaQBAHCKIQ0AgFMMaQAAnGJIAwDgFO3uBNR2dKptq7a3CyGEevXqlfdbgmOq5a+2Q4w1S9WWsUeOHDGZasqOHTtWrjlu3DiTFRUVmSxVz9zFqVu6dKnMb7/9dpP98pe/NFmXLl1Mdu+998o1Dx8+bDJ19rP6bKjrPYQQpk2bZjJ1vavPUCrgThoAAKcY0gAAOMWQBgDAKYY0AABOURxLYNSoUSZT24LGikCqsKCKZ6labMA3qetg+fLlJnv11Vfl8zds2GAyVZqZOnWqyWbPni3X3Ldvn8k40xwhxK+D//znPyZbuXKlyQYPHmyy66+/Xq7Zpk0bk23ZsiVR9sorr8g11Za3u3fvNlkqbAGqcCcNAIBTDGkAAJxiSAMA4BRDGgAApzJOJvy/6bHdtNKN+jnnzJljsvPOO89ksfOkBwwYYLJVq1aZrLi42GSpWnb4Lqf6M1WWa7A0Z5Xj1JTlc1VZrsPToUqVKomy0nwGSrObnzdJ3id30gAAOMWQBgDAKYY0AABOMaQBAHCKHce+Rf2P/CeeeMJkzz//vMkWLFgg1/ziiy8SvY7KMjMz5ZqqLIH0QEkM6UrtbsbOd9+NO2kAAJxiSAMA4BRDGgAApxjSAAA4xZAGAMAptgXFGce2oKhobAsKD9gWFACAFMaQBgDAKYY0AABOMaQBAHAqcXEMAACcWdxJAwDgFEMaAACnGNIAADjFkAYAwCmGNAAATjGkAQBwiiENAIBTDGkAAJxiSAMA4NT/AyYUf9Ev2JDMAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 500x500 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize some generated data, 3x3 grid\n",
    "fig, axs = plt.subplots(3,3, figsize=(5,5))\n",
    "axs = axs.flatten()\n",
    "for i in range(9):\n",
    "    axs[i].imshow(new_samples[i].view(28,28).detach().cpu().numpy(), cmap='gray')\n",
    "    axs[i].axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "Seems to work well for MNIST dataset. Tuning the hyperparameters might result in better images.\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}