diff --git a/generation/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb b/generation/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb
index b193e1023..67ecb5532 100644
--- a/generation/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb
+++ b/generation/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb
@@ -474,9 +474,8 @@
     }
    ],
    "source": [
-    "max_epochs = (\n",
-    "    1000  # training for longer (1e4 ~ 3h) helps a lot with reconstruction quality, even if the loss is already low\n",
-    ")\n",
+    "# training for longer (1e4 ~ 3h) helps a lot with reconstruction quality, even if the loss is already low\n",
+    "max_epochs = 1000\n",
     "val_interval = 100\n",
     "print_interval = 50\n",
     "iter_loss_list, val_iter_loss_list = [], []\n",