From e9f2d5809489d79f5a5c8043c39511ac77df1282 Mon Sep 17 00:00:00 2001 From: William Woof Date: Thu, 6 Apr 2023 18:13:24 +0100 Subject: [PATCH] Update onnx.py Fix break in computation graph due to typecast to python `int` instead of torch type. --- segment_anything/utils/onnx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/segment_anything/utils/onnx.py b/segment_anything/utils/onnx.py index 4297b31..493950a 100644 --- a/segment_anything/utils/onnx.py +++ b/segment_anything/utils/onnx.py @@ -81,8 +81,8 @@ class SamOnnxModel(nn.Module): 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]