diff --git a/sagemaker_batch_transform/working_with_tfrecords/working-with-tfrecords.ipynb b/sagemaker_batch_transform/working_with_tfrecords/working-with-tfrecords.ipynb
index ec20e480c9..ebe06ddfbe 100644
--- a/sagemaker_batch_transform/working_with_tfrecords/working-with-tfrecords.ipynb
+++ b/sagemaker_batch_transform/working_with_tfrecords/working-with-tfrecords.ipynb
@@ -50,7 +50,7 @@
     "import sagemaker\n",
     "import tensorflow as tf\n",
     "\n",
-    "bucket = \"<your_bucket_name>\"\n",
+    "bucket = sagemaker.Session().default_bucket()\n",
     "training_prefix = \"training\"\n",
     "batch_input_prefix = \"batch_input\"\n",
     "batch_output_prefix = \"batch_output\"\n",
@@ -292,7 +292,7 @@
     "tf_serving_model = Model(\n",
     "    model_data=estimator.model_data,\n",
     "    role=sagemaker.get_execution_role(),\n",
-    "    image=estimator.image_name,\n",
+    "    image_uri=estimator.image_uri,\n",
     "    framework_version=estimator.framework_version,\n",
     "    sagemaker_session=estimator.sagemaker_session,\n",
     ")\n",