Skip to content

GitLab

Unverified Commit e9f2d580 authored by William Woof's avatar William Woof Committed by GitHub
Browse files

Update onnx.py 

Fix break in computation graph due to typecast to python `int` instead of torch type.
parent f2557f77
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
segment_anything/utils/onnx.py
View file @ e9f2d580
......@@ -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]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment