fix get-cats function to add kwarg usage

NVIDIA-Merlin · Jun 15, 2023 · d00cb8f · d00cb8f
1 parent d26e776
commit d00cb8f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -279,7 +279,7 @@ def get_cats(workflow, col, stat_name="categories", cpu=False):
     # figure out the categorify node from the workflow graph
     cats = [
         cg.op
-        for cg in iter_nodes([workflow.output_node])
+        for cg in iter_nodes([workflow.output_node], flatten_subgraphs=True)
         if isinstance(cg.op, nvtabular.ops.Categorify)
     ]
     if len(cats) != 1: