From d00cb8fdb85ae7094b1cd0429e77ff7ad0afa5da Mon Sep 17 00:00:00 2001 From: Julio Date: Wed, 14 Jun 2023 21:42:23 -0400 Subject: [PATCH] fix get-cats function to add kwarg usage --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 3c3ae4373b..a5f5e1f6bb 100644 --- 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: