zjt/SAMapi
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update onnx.py

Browse Source 
Fix break in computation graph due to typecast to python `int` instead of torch type.
This commit is contained in:
William Woof
2023-04-06 18:13:24 +01:00
committed by GitHub
parent f2557f7780
commit e9f2d58094
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
segment_anything/utils/onnx.py
View File

@@ -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]