diff --git a/mnist_autoencoder.ipynb b/mnist_autoencoder.ipynb
index 52a58af..beed360 100644
--- a/mnist_autoencoder.ipynb
+++ b/mnist_autoencoder.ipynb
@@ -127,6 +127,37 @@
      "name": "#%%\n"
     }
    }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Autoencoder class"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "outputs": [],
+   "source": [
+    "class Autoencoder(nn.Module):\n",
+    "    def __init__(self, latent_dims):\n",
+    "        super(Autoencoder, self).__init__()\n",
+    "        self.encoder = Encoder(latent_dims)\n",
+    "        self.decoder = Decoder(latent_dims)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        z = self.encoder(x)\n",
+    "        return self.decoder(z)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
   }
  ],
  "metadata": {