diff --git a/examples/api/model_construction_tutorial.ipynb b/examples/api/model_construction_tutorial.ipynb
index e3d6236e..71330818 100644
--- a/examples/api/model_construction_tutorial.ipynb
+++ b/examples/api/model_construction_tutorial.ipynb
@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infer-actively/pymdp/blob/main/examples/api/model_construction_tutorial.ipynb)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -29,7 +36,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "if \"google.colab\" in sys.modules:\n",
+    "    %pip install \"inferactively-pymdp\" -q"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -44,7 +62,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -83,7 +101,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -122,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -133,6 +151,14 @@
       "Goal position: center_left\n",
       "Action chosen: move_right\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_1062531/1053406980.py:4: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n",
+      "  agent = Agent(**model, gamma=gamma)\n"
+     ]
     }
    ],
    "source": [
@@ -168,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -185,6 +211,14 @@
       "Agent 1's action chosen: move_left\n",
       "Agent 2's action chosen: move_left\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_1062531/1678643109.py:5: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n",
+      "  agent = Agent(**model, batch_size=batch_size, gamma=gamma)\n"
+     ]
     }
    ],
    "source": [
@@ -227,7 +261,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -294,7 +328,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -404,9 +438,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_1062531/53738707.py:3: UserWarning: A JAX array is being set as static! This can result in unexpected behavior and is usually a mistake to do.\n",
+      "  agent = Agent(**model, learn_A=False, learn_B=False, gamma=gamma, sampling_mode=\"full\")\n"
+     ]
+    }
+   ],
    "source": [
     "gamma = 1.0\n",
     "\n",
@@ -422,7 +465,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -440,7 +483,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -458,7 +501,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -469,7 +512,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -478,7 +521,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -498,7 +541,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -519,7 +562,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -527,25 +570,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv311",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.12"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 2
 }