diff --git a/examples_and_demos/kornia_image_points_transforms.ipynb b/examples_and_demos/kornia_image_points_transforms.ipynb
index 5c2fc157..21bedaf9 100644
--- a/examples_and_demos/kornia_image_points_transforms.ipynb
+++ b/examples_and_demos/kornia_image_points_transforms.ipynb
@@ -42,6 +42,7 @@
         "%%capture\n",
         "!pip install kornia\n",
         "!pip install kornia-rs matplotlib\n",
+        "!pip install flax",
         "!pip install -q ivy \n",
         "\n",
         "!rm -rf ivy_transpiled_outputs/"