diff --git a/examples/classical/demo.ipynb b/examples/classical/demo.ipynb
index 6c739e123..0de988b5b 100644
--- a/examples/classical/demo.ipynb
+++ b/examples/classical/demo.ipynb
@@ -134,7 +134,8 @@
     "pairs = nbList.pairs\n",
     "nbfunc = potentials[-1]\n",
     "energy = nbfunc(positions, box, pairs, ff.getGenerators()[-1].params)\n",
-    "print(energy)"
+    "print(energy)\n",
+    "print(pairs)"
    ]
   },
   {
@@ -198,9 +199,9 @@
    "hash": "44fe82502fda871be637af1aa98d2b3ddaac01204dd30f1519cbec4e95000815"
   },
   "kernelspec": {
-   "display_name": "dmff",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "dmff"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -212,7 +213,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.7.12"
   }
  },
  "nbformat": 4,
diff --git a/examples/fluctuated_leading_term_waterff/demo.ipynb b/examples/fluctuated_leading_term_waterff/demo.ipynb
index abc0b5982..8453e1da0 100644
--- a/examples/fluctuated_leading_term_waterff/demo.ipynb
+++ b/examples/fluctuated_leading_term_waterff/demo.ipynb
@@ -196,7 +196,7 @@
     "            dScales\n",
     "        )\n",
     "        return E_pme \n",
-    "    return admp_calculator\n"
+    "    return jax.jit(admp_calculator)\n"
    ]
   },
   {
@@ -231,20 +231,30 @@
     "print(force)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "c74539d4",
+   "metadata": {},
+   "source": [
+    "The running speed of the first pass is slow because JAX is trying to track the data flow and compile the code. Once the code is compiled, it is run much faster, until the shapes of the input parameters change, trigerring a recompilation.  "
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "d7a0fd1b-a363-40a0-9017-febe0fe76f01",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "print(-force_fn(positions, box, pairs))"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "dmff",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "dmff"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -256,7 +266,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.7.12"
   }
  },
  "nbformat": 4,