From 9d446d90fa8e059a5f46580912abd7e61cb414de Mon Sep 17 00:00:00 2001
From: SUSHMANTH REDDY <73489688+sushmanthreddy@users.noreply.github.com>
Date: Fri, 29 Sep 2023 18:23:19 +0530
Subject: [PATCH 1/2] Update mask_decoder.py

---
 segment_anything/modeling/mask_decoder.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/segment_anything/modeling/mask_decoder.py b/segment_anything/modeling/mask_decoder.py
index 5d2fdb03d..94e8a1e74 100644
--- a/segment_anything/modeling/mask_decoder.py
+++ b/segment_anything/modeling/mask_decoder.py
@@ -123,7 +123,12 @@ def predict_masks(
         tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
 
         # Expand per-image data in batch direction to be per-mask
-        src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
+        if image_embeddings.shape[0] != tokens.shape[0]:
+            src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
+        else:
+            src = image_embeddings
+
+        
         src = src + dense_prompt_embeddings
         pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
         b, c, h, w = src.shape

From c779d451d22743dc56a30ab7e7896fed3bfeb5fb Mon Sep 17 00:00:00 2001
From: SUSHMANTH REDDY <73489688+sushmanthreddy@users.noreply.github.com>
Date: Fri, 29 Sep 2023 19:53:56 +0530
Subject: [PATCH 2/2] Update mask_decoder.py

---
 segment_anything/modeling/mask_decoder.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/segment_anything/modeling/mask_decoder.py b/segment_anything/modeling/mask_decoder.py
index 94e8a1e74..9f799a5ff 100644
--- a/segment_anything/modeling/mask_decoder.py
+++ b/segment_anything/modeling/mask_decoder.py
@@ -128,7 +128,6 @@ def predict_masks(
         else:
             src = image_embeddings
 
-        
         src = src + dense_prompt_embeddings
         pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
         b, c, h, w = src.shape