Apply the expand to the dataframe instead.

Author: MarkDaoust

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",
@@ -502,7 +503,7 @@
         "test_type_layer = get_category_encoding_layer(name='Type',\n",
         "                                              dataset=train_ds,\n",
         "                                              dtype='string')\n",
-        "test_type_layer(tf.expand_dims(test_type_col, -1))"
+        "test_type_layer(test_type_col)"
       ]
     },
     {
@@ -527,7 +528,7 @@
         "                                             dataset=train_ds,\n",
         "                                             dtype='int64',\n",
         "                                             max_tokens=5)\n",
-        "test_age_layer(tf.expand_dims(test_age_col, -1))"
+        "test_age_layer(test_age_col)"
       ]
     },
     {