diff --git a/generation/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb b/generation/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb
index 4e1c9e1a6..72859b0f6 100644
--- a/generation/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb
+++ b/generation/3d_autoencoderkl/3d_autoencoderkl_tutorial.ipynb
@@ -694,6 +694,8 @@
     "n_evaluations = min(5, len(intermediary_images))\n",
     "\n",
     "fig, axs = plt.subplots(nrows=n_evaluations, ncols=3, constrained_layout=True, figsize=(8, 6))\n",
+    "if not isinstance(axs[0], tuple):\n",
+    "    axs = (axs,)\n",
     "\n",
     "# Remove ticks\n",
     "for ax in axs.flatten():\n",