Merge pull request facebookresearch#58 from EndingCredits/main

Fix baking original input size values to ONNX graph
phygitalism · Apr 9, 2023 · f58f1d1 · f58f1d1
2 parents efeab72 + e9f2d58
commit f58f1d1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/segment_anything/utils/onnx.py b/segment_anything/utils/onnx.py
@@ -81,8 +81,8 @@ def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -
             align_corners=False,
         )
 
-        prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size)
-        masks = masks[..., : int(prepadded_size[0]), : int(prepadded_size[1])]
+        prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size).to(torch.int64)
+        masks = masks[..., : prepadded_size[0], : prepadded_size[1]]
 
         orig_im_size = orig_im_size.to(torch.int64)
         h, w = orig_im_size[0], orig_im_size[1]