diff --git a/colab/edge_detection.ipynb b/colab/edge_detection.ipynb
index d8e325b..802d40a 100644
--- a/colab/edge_detection.ipynb
+++ b/colab/edge_detection.ipynb
@@ -61,10 +61,10 @@
       "metadata": {
         "id": "6YGqN2uqP_7P",
         "colab_type": "code",
-        "outputId": "66dfb8e7-f139-45cd-88c4-dbdf19d0c01e",
+        "outputId": "6a5f982c-17d9-485b-f95c-9656a1583c9b",
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 343
+          "height": 326
         }
       },
       "source": [
@@ -117,7 +117,7 @@
             "INFO:tensorflow:Assets written to: /usr/local/google/home/scotttodd/saved_models/edge_detection.sm/assets\n",
             "Edge Detection MLIR: \n",
             "\n",
-            "module {\n",
+            "module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 175 : i32}} {\n",
             "  func @edge_detect_sobel_operator(%arg0: tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> attributes {iree.module.export, iree.reflection = {abi = \"sip\", abiv = 1 : i32, sip = \"I8!S5!k0_0R3!_0\"}, tf._input_shapes = [\"tfshape$dim { size: 1 } dim { size: 128 } dim { size: 128 } dim { size: 1 }\"]} {\n",
             "    %0 = xla_hlo.constant dense<[[[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]], [[[-2.000000e+00]], [[0.000000e+00]], [[2.000000e+00]]], [[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n",
             "    %1 = xla_hlo.constant dense<[[[[1.000000e+00]], [[2.000000e+00]], [[1.000000e+00]]], [[[0.000000e+00]], [[0.000000e+00]], [[0.000000e+00]]], [[[-1.000000e+00]], [[-2.000000e+00]], [[-1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n",
@@ -130,7 +130,6 @@
             "    return %7 : tensor<1x128x128x1xf32>\n",
             "  }\n",
             "}\n",
-            "\n",
             "Wrote MLIR to path '/usr/local/google/home/scotttodd/saved_models/edge_detection.mlir'\n"
           ],
           "name": "stdout"
@@ -142,11 +141,11 @@
       "metadata": {
         "id": "Ytvb5Gx_EFJl",
         "colab_type": "code",
+        "outputId": "0f46bfc7-ebe6-4eb1-c65a-fc8b65adcada",
         "colab": {
           "base_uri": "https://localhost:8080/",
           "height": 51
-        },
-        "outputId": "6ba61508-3799-49dd-dd54-7c2b275764cf"
+        }
       },
       "source": [
         "#@title Prepare to test the edge detection module\n",
@@ -167,8 +166,8 @@
         {
           "output_type": "stream",
           "text": [
-            "Created IREE driver vulkan: <pyiree.rt.binding.HalDriver object at 0x7fd40c31b7b0>\n",
-            "SystemContext driver=<pyiree.rt.binding.HalDriver object at 0x7fd40c31b7b0>\n"
+            "Created IREE driver vulkan: <pyiree.rt.binding.HalDriver object at 0x7ff712a0f0b0>\n",
+            "SystemContext driver=<pyiree.rt.binding.HalDriver object at 0x7ff712a0f0b0>\n"
           ],
           "name": "stderr"
         }
@@ -205,7 +204,7 @@
       "metadata": {
         "id": "qW32e6spORCo",
         "colab_type": "code",
-        "outputId": "9699fc94-82b7-4ea8-fc72-45681ee0d3a3",
+        "outputId": "7b0070f3-85ac-46d4-e3f4-e990c5b7c501",
         "colab": {
           "base_uri": "https://localhost:8080/",
           "height": 570
@@ -222,13 +221,13 @@
         "\n",
         "# Plot the input and output images\n",
         "print(\"Input:\")\n",
-        "plt.imshow(arg0.reshape(128, 128), cmap=\"gray\")\n",
+        "plt.imshow(content_image.numpy().reshape(128, 128), cmap=\"gray\")\n",
         "plt.show()\n",
         "print(\"Output:\")\n",
         "plt.imshow(result.reshape(128, 128), cmap=\"gray\")\n",
         "plt.show()"
       ],
-      "execution_count": 6,
+      "execution_count": 5,
       "outputs": [
         {
           "output_type": "stream",