diff --git a/docs/tutorial_advanced/operator_custom_with_taichi.ipynb b/docs/tutorial_advanced/operator_custom_with_taichi.ipynb
index 70a88eeea..183a8a251 100644
--- a/docs/tutorial_advanced/operator_custom_with_taichi.ipynb
+++ b/docs/tutorial_advanced/operator_custom_with_taichi.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Operator Customization with Numba"
+    "# Operator Customization with Taichi"
    ]
   },
   {
@@ -18,7 +18,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Brain dynamics is sparse and event-driven, however, proprietary operators for brain dynamics are not well abstracted and summarized. As a result, we are often faced with the need to customize operators. In this tutorial, we will explore how to customize brain dynamics operators using ti.\n",
+    "Brain dynamics is sparse and event-driven, however, proprietary operators for brain dynamics are not well abstracted and summarized. As a result, we are often faced with the need to customize operators. In this tutorial, we will explore how to customize brain dynamics operators using taichi.\n",
     "\n",
     "Start by importing the relevant Python package."
    ]
@@ -242,11 +242,11 @@
     "                  vector: ti.types.ndarray(ndim=1), \n",
     "                  weight: ti.types.ndarray(ndim=1), \n",
     "                  out: ti.types.ndarray(ndim=1)):\n",
-    "   weight_0 = weight[0]\n",
-    "   ti.loop_config(block_dim=64)\n",
-    "   for i, j in ti.ndrange(indices.shape[0], indices.shape[1]):\n",
-    "        if vector[i]:\n",
-    "            out[indices[i, j]] += weight_0\n",
+    "  weight_0 = weight[0]\n",
+    "  ti.loop_config(block_dim=64)\n",
+    "  for ij in ti.grouped(indices):\n",
+    "      if vector[ij[0]]:\n",
+    "          out[ij[1]] += weight_0\n",
     "\n",
     "prim = bm.XLACustomOp(cpu_kernel=event_ell_cpu, gpu_kernel=event_ell_gpu)\n",
     "\n",
@@ -503,11 +503,11 @@
     "                  vector: ti.types.ndarray(ndim=1), \n",
     "                  weight: ti.types.ndarray(ndim=1), \n",
     "                  out: ti.types.ndarray(ndim=1)):\n",
-    "   weight_0 = weight[0]\n",
-    "   ti.loop_config(block_dim=64)\n",
-    "   for i, j in ti.ndrange(indices.shape[0], indices.shape[1]):\n",
-    "        if vector[i]:\n",
-    "            out[indices[i, j]] += weight_0\n",
+    "  weight_0 = weight[0]\n",
+    "  ti.loop_config(block_dim=64)\n",
+    "  for ij in ti.grouped(indices):\n",
+    "      if vector[ij[0]]:\n",
+    "          out[ij[1]] += weight_0\n",
     "\n",
     "prim = bm.XLACustomOp(cpu_kernel=event_ell_cpu, gpu_kernel=event_ell_gpu)\n",
     "\n",