Switching to categorical_crossentropy

PiperOrigin-RevId: 234160929

Author: yashk2810

site/en/r2/tutorials/distribute/keras.ipynb

@@ -190,7 +190,7 @@
         "num_train_examples = ds_info.splits['train'].num_examples\n",
         "num_test_examples = ds_info.splits['test'].num_examples\n",
         "\n",
-        "BUFFER_SIZE = num_train_examples\n",
+        "BUFFER_SIZE = 10000\n",
         "BATCH_SIZE = 64"
       ]
     },
@@ -217,7 +217,8 @@
         "def scale(image, label):\n",
         "  image = tf.cast(image, tf.float32)\n",
         "  image /= 255\n",
-        "  return image, label"
+        "    \n",
+        "  return image, label[..., tf.newaxis]"
       ]
     },
     {
@@ -318,10 +319,10 @@
         "      tf.keras.layers.Dense(64, activation='relu'),\n",
         "      tf.keras.layers.Dense(10, activation='softmax')\n",
         "  ])\n",
-        "  # TODO(yashkatariya): Add accuracy when b/122371345 is fixed.\n",
+        "\n",
         "  model.compile(loss='sparse_categorical_crossentropy',\n",
-        "                optimizer=tf.keras.optimizers.Adam())\n",
-        "                #metrics=['accuracy'])"
+        "                optimizer=tf.keras.optimizers.Adam(),\n",
+        "                metrics=['accuracy'])"
       ]
     },
     {