Skip to content

Commit c5a1826

Browse files
committed
add torch._check hints for torch.export
1 parent 536dc9d commit c5a1826

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

detectron2/structures/image_list.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,12 @@ def from_tensors(
111111
# This seems slightly (2%) faster.
112112
# TODO: check whether it's faster for multiple images as well
113113
image_size = image_sizes[0]
114-
padding_size = [0, max_size[-1] - image_size[1], 0, max_size[-2] - image_size[0]]
114+
u0 = (max_size[-1] - image_size[1]).item()
115+
u1 = (max_size[-2] - image_size[0]).item()
116+
padding_size = [0, u0, 0, u1]
117+
if torch.compiler.is_compiling():
118+
torch._check(u0 >= 0)
119+
torch._check(u1 >= 0)
115120
batched_imgs = F.pad(tensors[0], padding_size, value=pad_value).unsqueeze_(0)
116121
else:
117122
# max_size can be a tensor in tracing mode, therefore convert to list

0 commit comments

Comments
 (0)