diff --git a/main.ipynb b/main.ipynb
index ce4168b92d71f40bf56ec747f3abd60d74605e1d..ae8d2a25978dddb4fb5b1968833321e6f0d0feb9 100644
--- a/main.ipynb
+++ b/main.ipynb
@@ -295,7 +295,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "NAME = 'Baseline_VAE_short_512_10=warmup_begin'\n",
+    "NAME = 'Baseline_VAEP_short_512_10=warmup_begin'\n",
     "log_dir = os.path.abspath(\n",
     "    \"logs/fit/{}_{}\".format(NAME, datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
     ")\n",
@@ -326,10 +326,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "encoder, generator, vae, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAE(pttest.shape,\n",
-    "                                                                   loss='ELBO+MMD',\n",
-    "                                                                   kl_warmup_epochs=10,\n",
-    "                                                                   mmd_warmup_epochs=10).build()"
+    "# encoder, generator, vae, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAE(pttest.shape,\n",
+    "#                                                                    loss='ELBO+MMD',\n",
+    "#                                                                    kl_warmup_epochs=10,\n",
+    "#                                                                    mmd_warmup_epochs=10).build()"
    ]
   },
   {
@@ -338,10 +338,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# encoder, generator, vaep, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAEP(pttest.shape,\n",
-    "#                                                                     loss='ELBO+MMD',\n",
-    "#                                                                     kl_warmup_epochs=10,\n",
-    "#                                                                     mmd_warmup_epochs=10).build()"
+    "encoder, generator, vaep, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAEP(pttest.shape,\n",
+    "                                                                    loss='ELBO+MMD',\n",
+    "                                                                    kl_warmup_epochs=10,\n",
+    "                                                                    mmd_warmup_epochs=10).build()"
    ]
   },
   {
@@ -405,10 +405,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tf.config.experimental_run_functions_eagerly(False)\n",
-    "history = vae.fit(x=pttrain[:-1], y=pttrain[:-1], epochs=500, batch_size=512, verbose=1,\n",
-    "                  validation_data=(pttest[:-1], pttest[:-1]),\n",
-    "                  callbacks=[tensorboard_callback, kl_warmup_callback, mmd_warmup_callback])"
+    "# tf.config.experimental_run_functions_eagerly(False)\n",
+    "# history = vae.fit(x=pttrain[:-1], y=pttrain[:-1], epochs=500, batch_size=512, verbose=1,\n",
+    "#                   validation_data=(pttest[:-1], pttest[:-1]),\n",
+    "#                   callbacks=[tensorboard_callback, kl_warmup_callback, mmd_warmup_callback])"
    ]
   },
   {
@@ -419,10 +419,10 @@
    },
    "outputs": [],
    "source": [
-    "# tf.config.experimental_run_functions_eagerly(False)\n",
-    "# history = vaep.fit(x=pttrain[:-1], y=[pttrain[:-1],pttrain[1:]], epochs=500, batch_size=512, verbose=1,\n",
-    "#                    validation_data=(pttest[:-1], [pttest[:-1],pttest[1:]]),\n",
-    "#                    callbacks=[tensorboard_callback])"
+    "tf.config.experimental_run_functions_eagerly(False)\n",
+    "history = vaep.fit(x=pttrain[:-1], y=[pttrain[:-1],pttrain[1:]], epochs=500, batch_size=512, verbose=1,\n",
+    "                   validation_data=(pttest[:-1], [pttest[:-1],pttest[1:]]),\n",
+    "                   callbacks=[tensorboard_callback])"
    ]
   }
  ],