diff --git a/site/en/tutorials/distribute/keras.ipynb b/site/en/tutorials/distribute/keras.ipynb
index 2aa3f0a45f..b96656d443 100644
--- a/site/en/tutorials/distribute/keras.ipynb
+++ b/site/en/tutorials/distribute/keras.ipynb
@@ -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",