diff --git a/af2bind.ipynb b/af2bind.ipynb
index 8263f435..497ec592 100644
--- a/af2bind.ipynb
+++ b/af2bind.ipynb
@@ -3,8 +3,8 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/sokrypton/af2bind/blob/main/af2bind.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -123,7 +123,7 @@
         "  with open(f\"af2bind_params/attempt_7_2k_lam0-03/{model_type}.pickle\",\"rb\") as handle:\n",
         "    params_ = pickle.load(handle)\n",
         "  params_ = dict(**params_[\"~\"], **params_[\"linear\"])\n",
-        "  p = jax.tree_map(lambda x:np.asarray(x), params_)\n",
+        "  p = jax.tree.map(lambda x:np.asarray(x), params_)\n",
         "\n",
         "  # get predictions\n",
         "  x = (x - p[\"mean\"]) / p[\"std\"]\n",
@@ -135,6 +135,12 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "5r6hHIimktMH"
+      },
+      "outputs": [],
       "source": [
         "#@title **Run AF2BIND** 🔬\n",
         "target_pdb = \"6w70\" #@param {type:\"string\"}\n",
@@ -201,13 +207,7 @@
         "\n",
         "print(\"\\n🧪Pymol Selection Cmd:\")\n",
         "print(pymol_cmd)"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "5r6hHIimktMH"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
@@ -327,6 +327,12 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6LNbani6bCiz"
+      },
+      "outputs": [],
       "source": [
         "#@title **Activation analysis** (optional)\n",
         "pbind_cutoff = 0.5 # @param [\"0.0\", \"0.5\", \"0.9\"] {type:\"raw\"}\n",
@@ -350,20 +356,14 @@
         "                color_continuous_scale=[\"red\", \"white\", \"blue\"],\n",
         "              )\n",
         "fig.show()"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "6LNbani6bCiz"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
   ],
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "provenance": [],
-      "include_colab_link": true
+      "include_colab_link": true,
+      "provenance": []
     },
     "kernelspec": {
       "display_name": "Python 3",
@@ -375,4 +375,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}
diff --git a/af2bind_experimental.ipynb b/af2bind_experimental.ipynb
index 8fb64782..3f13e112 100644
--- a/af2bind_experimental.ipynb
+++ b/af2bind_experimental.ipynb
@@ -3,8 +3,8 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/sokrypton/af2bind/blob/main/af2bind_experimental.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -118,7 +118,7 @@
         "  with open(f\"af2bind_params/attempt_7_2k_lam0-03/{model_type}.pickle\",\"rb\") as handle:\n",
         "    params_ = pickle.load(handle)\n",
         "  params_ = dict(**params_[\"~\"], **params_[\"linear\"])\n",
-        "  p = jax.tree_map(lambda x:np.asarray(x), params_)\n",
+        "  p = jax.tree.map(lambda x:np.asarray(x), params_)\n",
         "\n",
         "  # get predictions\n",
         "  x = (x - p[\"mean\"]) / p[\"std\"]\n",
@@ -130,6 +130,12 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "5r6hHIimktMH"
+      },
+      "outputs": [],
       "source": [
         "#@title **Run AF2BIND** 🔬\n",
         "target_pdb = \"6w70\" #@param {type:\"string\"}\n",
@@ -227,13 +233,7 @@
         "\n",
         "print(\"\\n🧪Pymol Selection Cmd:\")\n",
         "print(pymol_cmd)"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "5r6hHIimktMH"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
@@ -353,6 +353,12 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6LNbani6bCiz"
+      },
+      "outputs": [],
       "source": [
         "#@title **Activation analysis** (optional)\n",
         "pbind_cutoff = 0.5 # @param [\"0.0\", \"0.5\", \"0.9\"] {type:\"raw\"}\n",
@@ -376,20 +382,14 @@
         "                color_continuous_scale=[\"red\", \"white\", \"blue\"],\n",
         "              )\n",
         "fig.show()"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "6LNbani6bCiz"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
   ],
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "provenance": [],
-      "include_colab_link": true
+      "include_colab_link": true,
+      "provenance": []
     },
     "kernelspec": {
       "display_name": "Python 3",
@@ -401,4 +401,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}
diff --git a/af2bind_large_pdb.ipynb b/af2bind_large_pdb.ipynb
index b4ff1870..2f322a01 100644
--- a/af2bind_large_pdb.ipynb
+++ b/af2bind_large_pdb.ipynb
@@ -128,7 +128,7 @@
         "  with open(f\"af2bind_params/attempt_7_2k_lam0-03/{model_type}.pickle\",\"rb\") as handle:\n",
         "    params_ = pickle.load(handle)\n",
         "  params_ = dict(**params_[\"~\"], **params_[\"linear\"])\n",
-        "  p = jax.tree_map(lambda x:np.asarray(x), params_)\n",
+        "  p = jax.tree.map(lambda x:np.asarray(x), params_) # Changed jax.tree_map to jax.tree.map\n",
         "\n",
         "  # get predictions\n",
         "  x = (x - p[\"mean\"]) / p[\"std\"]\n",