minor syntax fixes in integration tests (#351)

NVIDIA-Merlin · May 22, 2023 · a73cddd · a73cddd
1 parent 239cdad
commit a73cddd
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb b/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb
@@ -1026,8 +1026,8 @@
     "\n",
     "# read in data for request\n",
     "batch = df_lib.read_parquet(\n",
-    "    os.path.join(original_data_path,\"valid\", \"part.0.parquet\"), num_rows=3, columns=workflow.input_schema.column_names\n",
-    ")\n",
+    "    os.path.join(original_data_path,\"valid\", \"part.0.parquet\"), columns=workflow.input_schema.column_names\n",
+    ").head(3)\n",
     "batch"
    ]
   },

diff --git a/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py b/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py
@@ -1,7 +1,6 @@
 import os
 
 import pytest
-
 from testbook import testbook
 
 from tests.conftest import REPO_ROOT
@@ -97,9 +96,8 @@ def test_example_04_exporting_ranking_models(tb):
     # read in data for request
     batch = df_lib.read_parquet(
         os.path.join("/tmp/data/", "valid", "part.0.parquet"),
-        num_rows=3,
         columns=workflow.input_schema.column_names,
-    )
+    ).head(3)
     batch = batch.drop(columns="click")
     outputs = tb.ref("output_cols")
     from merlin.dataloader.tf_utils import configure_tensorflow

diff --git a/tests/integration/t4r/test_pytorch_backend.py b/tests/integration/t4r/test_pytorch_backend.py
@@ -34,7 +34,6 @@
 
 
 def test_serve_t4r_with_torchscript(tmpdir):
-
     # ===========================================
     # Generate training data
     # ===========================================
@@ -107,7 +106,7 @@ def test_serve_t4r_with_torchscript(tmpdir):
         if name in input_schema.column_names:
             dtype = input_schema[name].dtype
 
-            df_cols[name] = tensor.cpu().numpy().astype(dtype)
+            df_cols[name] = tensor.cpu().numpy().astype(dtype.name)
             if len(tensor.shape) > 1:
                 df_cols[name] = list(df_cols[name])