diff --git a/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb b/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb
index ddbb137d9..47c445180 100644
--- 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
index d16ab0a0c..2cdbde419 100644
--- 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
index 07d962c65..fe16f05a7 100644
--- 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])