Skip to content

Commit 65d50d0

Browse files
authored
Avoid creating a tensor of shape when not tracing (#7942)
Signed-off-by: Edward Z. Yang <[email protected]>
1 parent e13b8f5 commit 65d50d0

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

torchvision/models/detection/transform.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,10 @@ def _resize_image_and_masks(
3131
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
3232
if torchvision._is_tracing():
3333
im_shape = _get_shape_onnx(image)
34-
else:
34+
elif torch.jit.is_scripting():
3535
im_shape = torch.tensor(image.shape[-2:])
36+
else:
37+
im_shape = image.shape[-2:]
3638

3739
size: Optional[List[int]] = None
3840
scale_factor: Optional[float] = None

0 commit comments

Comments
 (0)