diff --git a/.DS_Store b/.DS_Store
index 6c999d7..d9688e6 100644
Binary files a/.DS_Store and b/.DS_Store differ
diff --git a/EfficientSAM_example.py b/EfficientSAM_example.py
index 3a86b09..997a270 100644
--- a/EfficientSAM_example.py
+++ b/EfficientSAM_example.py
@@ -8,14 +8,14 @@
 
 
 models = {}
-# Build the VIT-tiny model.
-models['vitt'] = build_efficient_sam_vitt()
+# Build the EfficientSAM-Ti model.
+models['efficientsam_ti'] = build_efficient_sam_vitt()
 
-# Since VIT-small is >100MB, we store the zip file.
+# Since EfficientSAM-S checkpoint is >100MB, we store the zip file.
 with zipfile.ZipFile("weights/efficient_sam_vits.pt.zip", 'r') as zip_ref:
     zip_ref.extractall("weights")
-# # Build the VIT-small model.
-models['vits'] = build_efficient_sam_vits()
+# # Build the EfficientSAM-S model.
+models['efficientsam_s'] = build_efficient_sam_vits()
 
 # load an image
 sample_image_np = np.array(Image.open("figs/examples/dogs.jpg"))
@@ -25,7 +25,7 @@
 input_points = torch.tensor([[[[580, 350], [650, 350]]]])
 input_labels = torch.tensor([[[1, 1]]])
 
-# Run inference for both vitt and vits based models.
+# Run inference for both EfficientSAM-Ti and EfficientSAM-S based models.
 for model_name, efficient_sam in models.items():
     print('Running inference using ', model_name)
     predicted_logits, predicted_iou = efficient_sam(
diff --git a/notebooks/EfficientSAM_example.ipynb b/notebooks/EfficientSAM_example.ipynb
index d4751da..8235717 100644
--- a/notebooks/EfficientSAM_example.ipynb
+++ b/notebooks/EfficientSAM_example.ipynb
@@ -17,7 +17,7 @@
         "id": "AIrAUKnLClPD"
       },
       "source": [
-        "This script provides example for how to get visualization result from EfficientSAM using ready-to-use torchscript, part of the code is borrow from MobileSAM project, many thanks!"
+        "This script provides example for how to get visualization result from EfficientSAM using checkpoint (ready-to-use torchscript can also be used), part of the code is borrow from MobileSAM project, many thanks!"
       ]
     },
     {
@@ -332,7 +332,6 @@
         }
       ],
       "source": [
-        "\n",
         "from efficient_sam import build_efficient_sam_vitt, build_efficient_sam_vits\n",
         "\n",
         "efficient_sam = build_efficient_sam_vitt()\n",
@@ -359,13 +358,6 @@
         "prepare your own image here"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": []
-    },
     {
       "cell_type": "code",
       "execution_count": 25,
@@ -487,19 +479,30 @@
     }
   ],
   "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "machine_shape": "hm",
-      "provenance": []
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
+    "custom": {
+      "cells": [],
+      "metadata": {
+        "accelerator": "GPU",
+        "colab": {
+          "machine_shape": "hm",
+          "provenance": []
+        },
+        "fileHeader": "",
+        "fileUid": "f337ddbb-4ec7-4bc4-8c8b-f31305249752",
+        "isAdHoc": false,
+        "kernelspec": {
+          "display_name": "Python 3",
+          "name": "python3"
+        },
+        "language_info": {
+          "name": "python"
+        }
+      },
+      "nbformat": 4,
+      "nbformat_minor": 0
     },
-    "language_info": {
-      "name": "python"
-    }
+    "indentAmount": 2
   },
   "nbformat": 4,
-  "nbformat_minor": 0
+  "nbformat_minor": 2
 }