From bad4aad0e4a8a85a4920bdc76b986e57c8529f50 Mon Sep 17 00:00:00 2001
From: Anna Grim <anna.grim@alleninstitute.org>
Date: Fri, 14 Aug 2026 02:30:27 +0000
Subject: [PATCH] refactor: pooled output from graph transformer

---
 src/arborist/data/datasets.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/arborist/data/datasets.py b/src/arborist/data/datasets.py
index b0f132e..2f16dd0 100644
--- a/src/arborist/data/datasets.py
+++ b/src/arborist/data/datasets.py
@@ -190,9 +190,10 @@ class TreeSample:
         branches.
     """
 
-    def __init__(self, curves, edge_index):
+    def __init__(self, curves, edge_index, root_curve_indices=None):
         self.curves = curves
         self.edge_index = edge_index
+        self.root_curve_indices = root_curve_indices
 
     def __repr__(self):
         return (