improve Categorify inference testing

NVIDIA-Merlin · Apr 10, 2024 · 337e904 · 337e904
1 parent 9f0ba33
commit 337e904
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/unit/ops/test_categorify.py b/tests/unit/ops/test_categorify.py
@@ -734,3 +734,8 @@ def test_categorify_inference():
     output_tensors = inference_op.transform(cats.input_columns, input_tensors)
     for key in input_tensors:
         assert output_tensors[key].dtype == np.dtype("int64")
+
+    # Check results are consistent with python code path
+    expect = workflow.transform(df)
+    got = pd.DataFrame(output_tensors)
+    assert_eq(expect, got)