Fix chapter 17 notebook

Author: mattdangerw

chapter17_image-generation.ipynb

@@ -534,16 +534,7 @@
     "\n",
     "    def compute_loss(self, x, y, y_pred, sample_weight=None, training=True):\n",
     "        _, pred_noise_masks, noise_masks = y_pred\n",
-    "        return self.loss(noise_masks, pred_noise_masks)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "colab_type": "text"
-   },
-   "source": [
-    "#### The generation process\n",
+    "        return self.loss(noise_masks, pred_noise_masks)\n",
     "\n",
     "    def generate(self, num_images, diffusion_steps):\n",
     "        noisy_images = keras.random.normal(\n",
@@ -570,6 +561,15 @@
     "        return ops.clip(images, 0.0, 255.0)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "#### The generation process"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {