diff --git a/site/en/tutorials/keras/save_and_load.ipynb b/site/en/tutorials/keras/save_and_load.ipynb
index 404fa1ee8b..bd417c095f 100644
--- a/site/en/tutorials/keras/save_and_load.ipynb
+++ b/site/en/tutorials/keras/save_and_load.ipynb
@@ -142,7 +142,7 @@
       },
       "outputs": [],
       "source": [
-        "!pip install pyyaml h5py  # Required to save models in HDF5 format"
+        "!pip install pyyaml h5py  # Required to save models in HDF5 format."
       ]
     },
     {
@@ -266,7 +266,7 @@
       },
       "outputs": [],
       "source": [
-        "checkpoint_path = \"training_1/cp.ckpt\"\n",
+        "checkpoint_path = \"training_1/cp.weights.h5\" # we are only saving weights therefore we need to use .weights.h5 extension instead we would use .keras for whole model\n",
         "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
         "\n",
         "# Create a callback that saves the model's weights\n",
@@ -275,8 +275,8 @@
         "                                                 verbose=1)\n",
         "\n",
         "# Train the model with the new callback\n",
-        "model.fit(train_images, \n",
-        "          train_labels,  \n",
+        "model.fit(train_images,\n",
+        "          train_labels,\n",
         "          epochs=10,\n",
         "          validation_data=(test_images, test_labels),\n",
         "          callbacks=[cp_callback])  # Pass callback to training\n",
@@ -312,7 +312,7 @@
         "id": "wlRN_f56Pqa9"
       },
       "source": [
-        "As long as two models share the same architecture you can share weights between them. So, when restoring a model from weights-only, create a model with the same architecture as the original model and then set its weights. \n",
+        "As long as two models share the same architecture you can share weights between them. So, when restoring a model from weights-only, create a model with the same architecture as the original model and then set its weights.\n",
         "\n",
         "Now rebuild a fresh, untrained model and evaluate it on the test set. An untrained model will perform at chance levels (~10% accuracy):"
       ]
@@ -380,8 +380,9 @@
       "outputs": [],
       "source": [
         "# Include the epoch in the file name (uses `str.format`)\n",
-        "checkpoint_path = \"training_2/cp-{epoch:04d}.ckpt\"\n",
+        "checkpoint_path = \"training_2/cp-{epoch:04d}.weights.h5\"\n",
         "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
+        "os.mkdir(checkpoint_dir)\n",
         "\n",
         "batch_size = 32\n",
         "\n",
@@ -392,8 +393,8 @@
         "\n",
         "# Create a callback that saves the model's weights every 5 epochs\n",
         "cp_callback = tf.keras.callbacks.ModelCheckpoint(\n",
-        "    filepath=checkpoint_path, \n",
-        "    verbose=1, \n",
+        "    filepath=checkpoint_path,\n",
+        "    verbose=1,\n",
         "    save_weights_only=True,\n",
         "    save_freq=5*n_batches)\n",
         "\n",
@@ -404,10 +405,10 @@
         "model.save_weights(checkpoint_path.format(epoch=0))\n",
         "\n",
         "# Train the model with the new callback\n",
-        "model.fit(train_images, \n",
+        "model.fit(train_images,\n",
         "          train_labels,\n",
-        "          epochs=50, \n",
-        "          batch_size=batch_size, \n",
+        "          epochs=50,\n",
+        "          batch_size=batch_size,\n",
         "          callbacks=[cp_callback],\n",
         "          validation_data=(test_images, test_labels),\n",
         "          verbose=0)"
@@ -441,7 +442,11 @@
       },
       "outputs": [],
       "source": [
-        "latest = tf.train.latest_checkpoint(checkpoint_dir)\n",
+        "def load_latest_checkpoint(checkpoint_dir):\n",
+        "    latest = max(os.listdir(checkpoint_dir), key=lambda f: int(f.split('-')[1].split('.')[0]))\n",
+        "    return os.path.join(checkpoint_dir, latest)\n",
+        "\n",
+        "latest = load_latest_checkpoint(checkpoint_dir)\n",
         "latest"
       ]
     },
@@ -505,7 +510,7 @@
       "source": [
         "## Manually save weights\n",
         "\n",
-        "To save weights manually, use `tf.keras.Model.save_weights`. By default, `tf.keras`—and the `Model.save_weights` method in particular—uses the TensorFlow [Checkpoint](../../guide/checkpoint.ipynb) format with a `.ckpt` extension. To save in the HDF5 format with a `.h5` extension, refer to the [Save and load models](https://www.tensorflow.org/guide/keras/save_and_serialize) guide."
+        "To save weights manually, use `tf.keras.Model.save_weights`. You have to use .weights.h5 extension to save the weights. You can refer to the [Save and load models](https://www.tensorflow.org/guide/keras/save_and_serialize) guide."
       ]
     },
     {
@@ -517,13 +522,14 @@
       "outputs": [],
       "source": [
         "# Save the weights\n",
-        "model.save_weights('./checkpoints/my_checkpoint')\n",
+        "os.mkdir('./checkpoints')\n",
+        "model.save_weights('./checkpoints/my_checkpoint.weights.h5')\n",
         "\n",
         "# Create a new model instance\n",
         "model = create_model()\n",
         "\n",
         "# Restore the weights\n",
-        "model.load_weights('./checkpoints/my_checkpoint')\n",
+        "model.load_weights('./checkpoints/my_checkpoint.weights.h5')\n",
         "\n",
         "# Evaluate the model\n",
         "loss, acc = model.evaluate(test_images, test_labels, verbose=2)\n",
@@ -657,7 +663,7 @@
         "id": "LtcN4VIb7JkK"
       },
       "source": [
-        "The SavedModel format is another way to serialize models. Models saved in this format can be restored using `tf.keras.models.load_model` and are compatible with TensorFlow Serving. The [SavedModel guide](../../guide/saved_model.ipynb) goes into detail about how to `serve/inspect` the SavedModel. The section below illustrates the steps to save and restore the model."
+        "The SavedModel format is another way to serialize models. Models saved in this format can directly be used with TFLite/TFServing/etc for inferencing. The [SavedModel guide](../../guide/saved_model.ipynb) goes into detail about how to `serve/inspect` the SavedModel. The section below illustrates the steps to save and restore the model."
       ]
     },
     {
@@ -674,7 +680,7 @@
         "\n",
         "# Save the entire model as a SavedModel.\n",
         "!mkdir -p saved_model\n",
-        "model.save('saved_model/my_model') "
+        "tf.saved_model.save(model, 'saved_model/my_model')"
       ]
     },
     {
@@ -707,7 +713,7 @@
         "id": "B7qfpvpY9HCe"
       },
       "source": [
-        "Reload a fresh Keras model from the saved model:"
+        "Reload the saved SavedModel file:"
       ]
     },
     {
@@ -718,34 +724,8 @@
       },
       "outputs": [],
       "source": [
-        "new_model = tf.keras.models.load_model('saved_model/my_model')\n",
-        "\n",
-        "# Check its architecture\n",
-        "new_model.summary()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "uWwgNaz19TH2"
-      },
-      "source": [
-        "The restored model is compiled with the same arguments as the original model. Try running evaluate and predict with the loaded model:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Yh5Mu0yOgE5J"
-      },
-      "outputs": [],
-      "source": [
-        "# Evaluate the restored model\n",
-        "loss, acc = new_model.evaluate(test_images, test_labels, verbose=2)\n",
-        "print('Restored model, accuracy: {:5.2f}%'.format(100 * acc))\n",
-        "\n",
-        "print(new_model.predict(test_images).shape)"
+        "saved_model = tf.saved_model.load('saved_model/my_model')\n",
+        "saved_model"
       ]
     },
     {
@@ -756,7 +736,7 @@
       "source": [
         "### HDF5 format\n",
         "\n",
-        "Keras provides a basic legacy high-level save format using the [HDF5](https://en.wikipedia.org/wiki/Hierarchical_Data_Format) standard. "
+        "Keras provides a basic legacy high-level save format using the [HDF5](https://en.wikipedia.org/wiki/Hierarchical_Data_Format) standard."
       ]
     },
     {
@@ -773,7 +753,7 @@
         "\n",
         "# Save the entire model to a HDF5 file.\n",
         "# The '.h5' extension indicates that the model should be saved to HDF5.\n",
-        "model.save('my_model.h5') "
+        "model.save('my_model.h5')"
       ]
     },
     {
@@ -859,21 +839,11 @@
         "\n",
         "Refer to the [Writing layers and models from scratch](https://www.tensorflow.org/guide/keras/custom_layers_and_models) tutorial for examples of custom objects and `get_config`.\n"
       ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "jBVTkkUIkEF3"
-      },
-      "outputs": [],
-      "source": []
     }
   ],
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "collapsed_sections": [],
       "name": "save_and_load.ipynb",
       "toc_visible": true
     },