diff --git a/micro_sam/instance_segmentation.py b/micro_sam/instance_segmentation.py
index d1324613..7e7a1107 100644
--- a/micro_sam/instance_segmentation.py
+++ b/micro_sam/instance_segmentation.py
@@ -907,7 +907,7 @@ def initialize(
         pbar_init(1, "Initialize instance segmentation with decoder")
 
         if image_embeddings is None:
-            image_embeddings = util.precompute_image_embeddings(self._predictor, image)
+            image_embeddings = util.precompute_image_embeddings(self._predictor, image, ndim=2)
 
         # Get the image embeddings from the predictor.
         self._predictor = util.set_precomputed(self._predictor, image_embeddings, i=i)