diff --git a/docs/source/notebooks/Example_4_Iowa_Gambling_Task_Short.ipynb b/docs/source/notebooks/Example_4_Iowa_Gambling_Task_Short.ipynb
index cecb3996f..0d3751a2e 100644
--- a/docs/source/notebooks/Example_4_Iowa_Gambling_Task_Short.ipynb
+++ b/docs/source/notebooks/Example_4_Iowa_Gambling_Task_Short.ipynb
@@ -38,12 +38,11 @@
     "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
     "import seaborn as sns\n",
-    "from jax import jit, vjp\n",
+    "from jax import jit\n",
     "from jax.nn import softmax\n",
     "from jax.tree_util import Partial\n",
-    "from pyhgf.math import binary_surprise, gaussian_surprise\n",
+    "from pyhgf.math import binary_surprise\n",
     "from pyhgf.model import HGF\n",
-    "from pyhgf.utils import beliefs_propagation\n",
     "from pytensor.graph import Apply, Op\n",
     "\n",
     "jax.config.update(\"jax_enable_x64\", True)  # this is required for the softmax\n",