Skip to content

Commit ec52805

Browse files
yushangdifacebook-github-bot
authored andcommitted
Add torch._check hints for torch.export
Summary: Fixes #5347 This allows for torch.export.export ImageList.from_tensors. Differential Revision: D74835454
1 parent 536dc9d commit ec52805

File tree

1 file changed

+19
-1
lines changed

1 file changed

+19
-1
lines changed

detectron2/structures/image_list.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,19 @@
88
from detectron2.layers.wrappers import move_device_like, shapes_to_tensor
99

1010

11+
try:
12+
try:
13+
from torch.compiler import is_compiling
14+
15+
except Exception:
16+
# torch.compiler.is_compiling is not available in torch 1.10
17+
from torch._dynamo import is_compiling as is_compiling
18+
19+
except Exception:
20+
21+
def is_compiling() -> bool: # type: ignore[misc]
22+
return False
23+
1124
class ImageList:
1225
"""
1326
Structure that holds a list of images (of possibly
@@ -111,7 +124,12 @@ def from_tensors(
111124
# This seems slightly (2%) faster.
112125
# TODO: check whether it's faster for multiple images as well
113126
image_size = image_sizes[0]
114-
padding_size = [0, max_size[-1] - image_size[1], 0, max_size[-2] - image_size[0]]
127+
u0 = max_size[-1] - image_size[1]
128+
u1 = max_size[-2] - image_size[0]
129+
padding_size = [0, u0, 0, u1]
130+
if is_compiling():
131+
torch._check(u0.item() >= 0)
132+
torch._check(u1.item() >= 0)
115133
batched_imgs = F.pad(tensors[0], padding_size, value=pad_value).unsqueeze_(0)
116134
else:
117135
# max_size can be a tensor in tracing mode, therefore convert to list

0 commit comments

Comments
 (0)