diff --git a/site/en/tutorials/structured_data/preprocessing_layers.ipynb b/site/en/tutorials/structured_data/preprocessing_layers.ipynb
index a94fdaf0c54..65762628eb5 100644
--- a/site/en/tutorials/structured_data/preprocessing_layers.ipynb
+++ b/site/en/tutorials/structured_data/preprocessing_layers.ipynb
@@ -295,9 +295,10 @@
       "outputs": [],
       "source": [
         "def df_to_dataset(dataframe, shuffle=True, batch_size=32):\n",
-        "  dataframe = dataframe.copy()\n",
-        "  labels = dataframe.pop('target')\n",
-        "  ds = tf.data.Dataset.from_tensor_slices((dict(dataframe), labels))\n",
+        "  df = dataframe.copy()\n",
+        "  labels = df.pop('target')\n",
+        "  df = {key: value[:,tf.newaxis] for key, value in dataframe.items()}\n",
+        "  ds = tf.data.Dataset.from_tensor_slices((dict(df), labels))\n",
         "  if shuffle:\n",
         "    ds = ds.shuffle(buffer_size=len(dataframe))\n",
         "  ds = ds.batch(batch_size)\n",