diff --git a/README.md b/README.md
index 6e87ca8..dcabf36 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ Grounded SAM 2 does not introduce significant methodological changes compared to
 
 ## Contents
 - [Installation](#installation)
-- [Grounded SAM 2 Demo](#grounded-sam-2-demo)
+- [Grounded SAM 2 Demo](#grounded-sam-2-demos)
   - [Grounded SAM 2 Image Demo](#grounded-sam-2-image-demo-with-grounding-dino)
   - [Grounded SAM 2 Image Demo (with Grounding DINO 1.5 & 1.6)](#grounded-sam-2-image-demo-with-grounding-dino-15--16)
   - [Grounded SAM 2 Video Object Tracking Demo](#grounded-sam-2-video-object-tracking-demo)
@@ -65,7 +65,7 @@ wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alp
 wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha2/groundingdino_swinb_cogcoor.pth
 ```
 
-## Grounded SAM 2 Demo
+## Grounded SAM 2 Demos
 ### Grounded SAM 2 Image Demo (with Grounding DINO)
 Note that `Grounding DINO` has already been supported in [Huggingface](https://huggingface.co/IDEA-Research/grounding-dino-tiny), so we provide two choices for running `Grounded SAM 2` model:
 - Use huggingface API to inference Grounding DINO (which is simple and clear)
diff --git a/grounded_sam2_gd1.5_demo.py b/grounded_sam2_gd1.5_demo.py
index f6ab623..95bbf1b 100644
--- a/grounded_sam2_gd1.5_demo.py
+++ b/grounded_sam2_gd1.5_demo.py
@@ -89,11 +89,7 @@ masks, scores, logits = sam2_predictor.predict(
 Post-process the output of the model to get the masks, scores, and logits for visualization
 """
 # convert the shape to (n, H, W)
-if masks.ndim == 3:
-    masks = masks[None]
-    scores = scores[None]
-    logits = logits[None]
-elif masks.ndim == 4:
+if masks.ndim == 4:
     masks = masks.squeeze(1)
 
 
diff --git a/grounded_sam2_hf_model_demo.py b/grounded_sam2_hf_model_demo.py
index b912e3f..10353fa 100644
--- a/grounded_sam2_hf_model_demo.py
+++ b/grounded_sam2_hf_model_demo.py
@@ -81,11 +81,7 @@ masks, scores, logits = sam2_predictor.predict(
 Post-process the output of the model to get the masks, scores, and logits for visualization
 """
 # convert the shape to (n, H, W)
-if masks.ndim == 3:
-    masks = masks[None]
-    scores = scores[None]
-    logits = logits[None]
-elif masks.ndim == 4:
+if masks.ndim == 4:
     masks = masks.squeeze(1)
 
 
diff --git a/grounded_sam2_local_demo.py b/grounded_sam2_local_demo.py
index f7e309f..2735e85 100644
--- a/grounded_sam2_local_demo.py
+++ b/grounded_sam2_local_demo.py
@@ -68,11 +68,7 @@ masks, scores, logits = sam2_predictor.predict(
 Post-process the output of the model to get the masks, scores, and logits for visualization
 """
 # convert the shape to (n, H, W)
-if masks.ndim == 3:
-    masks = masks[None]
-    scores = scores[None]
-    logits = logits[None]
-elif masks.ndim == 4:
+if masks.ndim == 4:
     masks = masks.squeeze(1)