diff --git a/samples/snippets/imported_onnx_model_test.py b/samples/snippets/imported_onnx_model_test.py
index a74d673087..87157ee60d 100644
--- a/samples/snippets/imported_onnx_model_test.py
+++ b/samples/snippets/imported_onnx_model_test.py
@@ -20,7 +20,7 @@ def test_imported_sklearn_onnx_model() -> None:
 
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev")
 
-    # [START bigquery_dataframes_imported_sklearn_onnx_tutorial_import_tensorflow_models]
+    # [START bigquery_dataframes_imported_sklearn_onnx_tutorial_import_onnx_models]
     import bigframes
     from bigframes.ml.imported import ONNXModel
 
@@ -31,7 +31,7 @@ def test_imported_sklearn_onnx_model() -> None:
     imported_onnx_model = ONNXModel(
         model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx"
     )
-    # [END bigquery_dataframes_imported_sklearn_onnx_tutorial_import_tensorflow_models]
+    # [END bigquery_dataframes_imported_sklearn_onnx_tutorial_import_onnx_models]
     assert imported_onnx_model is not None
 
     # [START bigquery_dataframes_imported_sklearn_onnx_tutorial_make_predictions]
@@ -39,5 +39,5 @@ def test_imported_sklearn_onnx_model() -> None:
 
     df = bpd.read_gbq("bigquery-public-data.ml_datasets.iris")
     predictions = imported_onnx_model.predict(df)
-    predictions.head(5)
+    predictions.peek(5)
     # [END bigquery_dataframes_imported_sklearn_onnx_tutorial_make_predictions]
diff --git a/samples/snippets/imported_tensorflow_model_test.py b/samples/snippets/imported_tensorflow_model_test.py
index 4913c635c2..7dcf0ffe6d 100644
--- a/samples/snippets/imported_tensorflow_model_test.py
+++ b/samples/snippets/imported_tensorflow_model_test.py
@@ -38,6 +38,7 @@ def test_imported_tensorflow_model() -> None:
     import bigframes.pandas as bpd
 
     df = bpd.read_gbq("bigquery-public-data.hacker_news.full")
-    predictions = imported_tensorflow_model.predict(df)
+    df_pred = df.rename(columns={"title": "input"})
+    predictions = imported_tensorflow_model.predict(df_pred)
     predictions.head(5)
     # [END bigquery_dataframes_imported_tensorflow_tutorial_make_predictions]