diff --git a/site/en/tutorials/distribute/keras.ipynb b/site/en/tutorials/distribute/keras.ipynb
index 1b6311a822c..b96656d4436 100644
--- a/site/en/tutorials/distribute/keras.ipynb
+++ b/site/en/tutorials/distribute/keras.ipynb
@@ -363,7 +363,7 @@
         "# Define the checkpoint directory to store the checkpoints.\n",
         "checkpoint_dir = './training_checkpoints'\n",
         "# Define the name of the checkpoint files.\n",
-        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}\")"
+        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch:04d}.weights.h5\")"
       ]
     },
     {
@@ -396,7 +396,7 @@
         "# Define a callback for printing the learning rate at the end of each epoch.\n",
         "class PrintLR(tf.keras.callbacks.Callback):\n",
         "  def on_epoch_end(self, epoch, logs=None):\n",
-        "    print('\\nLearning rate for epoch {} is {}'.format(        epoch + 1, model.optimizer.lr.numpy()))"
+        "    print('\\nLearning rate for epoch {} is {}'.format(epoch + 1, model.optimizer.learning_rate.numpy()))"
       ]
     },
     {
@@ -486,7 +486,10 @@
       },
       "outputs": [],
       "source": [
-        "model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))\n",
+        "import pathlib\n",
+        "latest_checkpoint = sorted(pathlib.Path(checkpoint_dir).glob('*'))[-1]\n",
+        "\n",
+        "model.load_weights(latest_checkpoint)\n",
         "\n",
         "eval_loss, eval_acc = model.evaluate(eval_dataset)\n",
         "\n",