diff --git a/examples/lats/lats.ipynb b/examples/lats/lats.ipynb
index bac0515b8..f514bf5c4 100644
--- a/examples/lats/lats.ipynb
+++ b/examples/lats/lats.ipynb
@@ -138,7 +138,8 @@
     "        \"\"\"Select the child with the highest UCT to search next.\"\"\"\n",
     "        if not self.children:\n",
     "            return None\n",
-    "        return max(self.children, key=lambda child: child.upper_confidence_bound())\n",
+    "        all_nodes = self._get_all_children()\n",
+    "        return max(all_nodes, key=lambda child: child.upper_confidence_bound())\n",
     "\n",
     "    @property\n",
     "    def best_child_score(self):\n",
@@ -190,9 +191,8 @@
     "            node = node.parent\n",
     "        # Reverse the final back-tracked trajectory to return in the correct order\n",
     "        return messages[::-1]  # root solution, reflection, child 1, ...\n",
-    "\n",
-    "    def get_best_solution(self):\n",
-    "        \"\"\"Return the best solution from within the current sub-tree.\"\"\"\n",
+    "    \n",
+    "    def _get_all_children(self):\n",
     "        all_nodes = [self]\n",
     "        nodes = deque()\n",
     "        nodes.append(self)\n",
@@ -201,6 +201,11 @@
     "            all_nodes.extend(node.children)\n",
     "            for n in node.children:\n",
     "                nodes.append(n)\n",
+    "        return all_nodes\n",
+    "\n",
+    "    def get_best_solution(self):\n",
+    "        \"\"\"Return the best solution from within the current sub-tree.\"\"\"\n",
+    "        all_nodes = self._get_all_children()\n",
     "        best_node = max(\n",
     "            all_nodes,\n",
     "            # We filter out all non-terminal, non-solution trajectories\n",