diff --git a/tests/aliby/network/extraction/test_base.py b/tests/aliby/network/extraction/test_base.py
index f1d079d27af7b8089bbb73ce2ca6b76f3dc54748..59ebc9766418f96c2a8bfbf37e27677dee498a45 100644
--- a/tests/aliby/network/extraction/test_base.py
+++ b/tests/aliby/network/extraction/test_base.py
@@ -52,7 +52,7 @@ def test_extractor(imgs, masks, tree):
     extractor.tree = tree
     traps = imgs["GFP"]
     # Generate mock labels
-    labels = list(range(masks.shape[2]))
+    labels = list(range(len(masks)))
     for ch_branches in extractor.params.tree.values():
         print(
             extractor.reduce_extract(