diff --git a/17_autoencoders_gans_and_diffusion_models.ipynb b/17_autoencoders_gans_and_diffusion_models.ipynb
index d8c727c..3d0bf33 100644
--- a/17_autoencoders_gans_and_diffusion_models.ipynb
+++ b/17_autoencoders_gans_and_diffusion_models.ipynb
@@ -83,6 +83,25 @@
     "assert version.parse(sklearn.__version__) >= version.parse(\"1.0.1\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Warning**: the latest TensorFlow versions are based on Keras 3. For chapters 10-15, it wasn't too hard to update the code to support Keras 3, but unfortunately it's much harder for this chapter, in particular adding custom losses using the functional API is not implemented yet. So for this chapter I've had to revert to Keras 2. To do that, I set the `TF_USE_LEGACY_KERAS` environment variable to `\"1\"` and import the `tf_keras` package. This ensures that `tf.keras` points to `tf_keras`, which is Keras 2.*."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\" \n",
+    "\n",
+    "import tf_keras"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -94,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {
     "id": "0Piq5se2pKzx"
    },
@@ -116,7 +135,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {
     "id": "8d4TH3NbpKzx"
    },
@@ -142,7 +161,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {
     "id": "PQFH5Y9PpKzy"
    },
@@ -171,7 +190,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "metadata": {
     "id": "Ekxzo6pOpKzy"
    },
@@ -202,7 +221,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -227,7 +246,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -248,7 +267,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -258,7 +277,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -309,7 +328,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -331,7 +350,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
@@ -420,7 +439,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -474,7 +493,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 15,
    "metadata": {
     "id": "g69nLh2GyLhC",
     "outputId": "47f84d4f-b41f-4ef9-d653-d73506cb4839"
@@ -499,7 +518,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 16,
    "metadata": {
     "id": "klppwl4gyLhC",
     "outputId": "b66330ea-5d49-4153-82ce-39fb78614a64"
@@ -534,7 +553,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 17,
    "metadata": {
     "id": "JIkuJSpPyLhC",
     "outputId": "239e4fe5-7756-40c1-cc2c-70c83eee6c0f"
@@ -598,7 +617,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 18,
    "metadata": {
     "id": "jTru0q0CyLhD"
    },
@@ -623,7 +642,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 19,
    "metadata": {
     "id": "f_mU3e6IyLhD",
     "outputId": "e8045bff-7004-4940-8d00-295ad99c0f23"
@@ -684,7 +703,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 20,
    "metadata": {
     "id": "WS2bm471yLhD",
     "outputId": "8427f9cc-630f-4cc9-a1f5-7361c5eab6dc",
@@ -721,7 +740,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 21,
    "metadata": {
     "id": "CIqJ922pyLhE"
    },
@@ -741,7 +760,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 22,
    "metadata": {
     "id": "-3TSUkv4yLhE",
     "outputId": "bca2fb96-7ddf-4457-b1df-30cdce8d803e"
@@ -807,7 +826,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 23,
    "metadata": {
     "id": "FpjUrD77yLhE"
    },
@@ -822,7 +841,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 24,
    "metadata": {
     "id": "KbEe4tQqyLhE",
     "outputId": "52eb1adf-f4c4-4029-f68b-1ef420a5daed"
@@ -855,7 +874,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 25,
    "metadata": {
     "id": "vRgEnA-WyLhF",
     "outputId": "3d367c4e-82d4-48d9-8757-b106edf00650"
@@ -886,7 +905,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 26,
    "metadata": {
     "id": "Ln1j0u4eyLhF",
     "outputId": "fef0981b-6b8d-4e9d-e837-9e3c20358e9e"
@@ -930,7 +949,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 27,
    "metadata": {
     "id": "qom8ZHdgyLhF",
     "outputId": "45b2544f-0812-4fc7-a4e5-d8df9f193f5f"
@@ -996,7 +1015,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 28,
    "metadata": {
     "id": "6p0Yud0RyLhG",
     "outputId": "7112e556-b714-4645-af12-3bb3bd7cd6ce"
@@ -1039,7 +1058,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 29,
    "metadata": {
     "id": "ZGOuVy0kyLhG"
    },
@@ -1062,7 +1081,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 30,
    "metadata": {
     "id": "4grialnWyLhG",
     "outputId": "e6bb8d79-9a57-4fd3-e336-b81687de6781"
@@ -1102,7 +1121,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 31,
    "metadata": {
     "id": "ZLOnTctFyLhG",
     "outputId": "e5b59e40-09bc-4d82-bb6d-58ae7d7cc376"
@@ -1146,7 +1165,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 32,
    "metadata": {
     "id": "wTKWjK3EyLhH",
     "outputId": "3be1a3bc-7f30-4793-955e-07bf60fc3b6e"
@@ -1203,7 +1222,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 33,
    "metadata": {
     "id": "Pucj7MI3yLhI",
     "outputId": "d9a4289a-fde7-4e3b-c968-72421ddbfe6d"
@@ -1258,7 +1277,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 34,
    "metadata": {
     "id": "ORrSHhNxyLhJ",
     "outputId": "664ca147-e8f2-4ccf-960e-a86e60c2c60d"
@@ -1315,7 +1334,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 34,
+   "execution_count": 35,
    "metadata": {
     "id": "2vU9twBAyLhK",
     "outputId": "feb78185-b67e-4775-c61e-3d4a8d082c75"
@@ -1351,7 +1370,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": 36,
    "metadata": {
     "id": "D0fOy7aMyLhK",
     "outputId": "35a609c2-6d23-462f-a519-6d93b07d0ec4"
@@ -1399,7 +1418,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 37,
    "metadata": {
     "id": "N-MoRrQGyLhK"
    },
@@ -1428,7 +1447,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 37,
+   "execution_count": 38,
    "metadata": {
     "id": "HIoQ-1D4yLhK",
     "outputId": "8874ead0-468a-47f8-9e35-4a855bace6b4"
@@ -1486,7 +1505,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 38,
+   "execution_count": 39,
    "metadata": {
     "id": "o0jIPg2ayLhL",
     "outputId": "35512003-0640-4b9a-835b-6c7af95c930f"
@@ -1522,7 +1541,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 39,
+   "execution_count": 40,
    "metadata": {
     "id": "64ji8G6eyLhL"
    },
@@ -1536,7 +1555,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": 41,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1557,7 +1576,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": 42,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1571,7 +1590,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": 43,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1582,7 +1601,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 43,
+   "execution_count": 44,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1594,7 +1613,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 44,
+   "execution_count": 45,
    "metadata": {
     "id": "YXUKW_UPyLhL",
     "outputId": "6db5b948-1b3d-419b-e810-08512fa65369"
@@ -1665,7 +1684,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 45,
+   "execution_count": 46,
    "metadata": {
     "id": "FDs525puyLhL",
     "outputId": "049794c0-4407-49b7-ba65-138f31d5e5dc",
@@ -1708,7 +1727,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 46,
+   "execution_count": 47,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1727,7 +1746,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 47,
+   "execution_count": 48,
    "metadata": {
     "id": "gzwLpDIDyLhM"
    },
@@ -1775,7 +1794,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 48,
+   "execution_count": 49,
    "metadata": {
     "id": "GsVXES9byLhM",
     "outputId": "eca5be95-9eee-47d1-e9ce-50756e24d235"
@@ -1791,7 +1810,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 49,
+   "execution_count": 50,
    "metadata": {},
    "outputs": [
     {
@@ -1825,7 +1844,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 51,
    "metadata": {
     "id": "N31KxgksyLhN"
    },
@@ -1853,7 +1872,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 52,
    "metadata": {
     "id": "Qm4Uk3MTyLhN"
    },
@@ -1866,7 +1885,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": 53,
    "metadata": {
     "id": "UQoPjJthyLhN"
    },
@@ -1879,7 +1898,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": 54,
    "metadata": {
     "id": "cxrUUhrjyLhN",
     "outputId": "2813bb7d-84c4-4ed2-8d03-c1fbd138c576"
@@ -2026,7 +2045,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": 55,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -2038,7 +2057,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 56,
    "metadata": {
     "id": "NrLzMxweyLhO",
     "outputId": "b21da95e-da35-45b6-bbee-fd6e8b915382"
@@ -2075,7 +2094,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 56,
+   "execution_count": 57,
    "metadata": {
     "id": "NgT77mdVyLhO"
    },
@@ -2110,7 +2129,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 58,
    "metadata": {
     "id": "0-Sj_zc9yLhO"
    },
@@ -2124,7 +2143,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 58,
+   "execution_count": 59,
    "metadata": {
     "id": "4Ss7zwLGyLhO"
    },
@@ -2135,7 +2154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 59,
+   "execution_count": 60,
    "metadata": {
     "id": "xKWHUiUqyLhP",
     "outputId": "576b04b2-5604-405a-a672-9693db6ef26a"
@@ -2268,7 +2287,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 60,
+   "execution_count": 61,
    "metadata": {
     "id": "yY8_xytVyLhP",
     "outputId": "7f3ebd45-3bc5-42ce-d0a6-5734d360cef6"
@@ -2318,7 +2337,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 61,
+   "execution_count": 62,
    "metadata": {
     "id": "yxPNTP3QpHAw"
    },
@@ -2358,7 +2377,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 62,
+   "execution_count": 63,
    "metadata": {},
    "outputs": [
     {
@@ -2401,7 +2420,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 63,
+   "execution_count": 64,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -2427,7 +2446,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 64,
+   "execution_count": 65,
    "metadata": {
     "id": "i9qt5HPZ8rhv"
    },
@@ -2455,7 +2474,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 65,
+   "execution_count": 66,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -2559,7 +2578,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 66,
+   "execution_count": 67,
    "metadata": {
     "id": "0HGLhS4oNkNC"
    },
@@ -2597,7 +2616,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 67,
+   "execution_count": 68,
    "metadata": {
     "id": "4QLnc8wkfXrh"
    },
@@ -2674,7 +2693,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 68,
+   "execution_count": 69,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -2764,7 +2783,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 69,
+   "execution_count": 70,
    "metadata": {
     "id": "vMXxz4qV8Luk"
    },
@@ -2797,7 +2816,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 70,
+   "execution_count": 71,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -2853,7 +2872,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 71,
+   "execution_count": 72,
    "metadata": {
     "id": "3DRvS7twyLhW",
     "outputId": "80a63cbf-0f33-4b6c-8e25-32474da2e2fb"
@@ -2917,7 +2936,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 72,
+   "execution_count": 73,
    "metadata": {
     "id": "HyOsfAdiyLhW",
     "outputId": "678cc611-c625-434e-90b9-ad90d88baef4"
@@ -2952,7 +2971,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 73,
+   "execution_count": 74,
    "metadata": {
     "id": "Akb6fNptyLhX",
     "outputId": "efdbf0bd-1506-42ae-f1e1-7cd68d0f45a7"
@@ -2991,7 +3010,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 74,
+   "execution_count": 75,
    "metadata": {
     "id": "hS80DfcLyLhX",
     "outputId": "51a38569-f02e-4737-d40e-b478d8a557b3"
@@ -3084,7 +3103,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 75,
+   "execution_count": 76,
    "metadata": {
     "id": "RVAVmIznyLhY"
    },
@@ -3097,7 +3116,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 76,
+   "execution_count": 77,
    "metadata": {
     "id": "qXcft-kHyLhY"
    },
@@ -3116,7 +3135,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 77,
+   "execution_count": 78,
    "metadata": {
     "id": "a2lC7XBqyLhY"
    },
@@ -3132,7 +3151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 78,
+   "execution_count": 79,
    "metadata": {
     "id": "RgtxA9BjyLhY",
     "outputId": "7b43827a-cc75-45cd-e5e0-dbd4ccfea440"
@@ -3175,7 +3194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 79,
+   "execution_count": 80,
    "metadata": {
     "id": "dTlFyrmSyLhY",
     "outputId": "7a842325-f99e-4e85-962f-ba3bcbf8cdfd"
@@ -3281,7 +3300,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.9.10"
   },
   "nav_menu": {
    "height": "381px",