From d50177f257e5193019c4f09031e21ef05659097c Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Thu, 7 Sep 2023 17:09:54 -0400 Subject: [PATCH] Avoid creating a tensor of shape when not tracing This is a small follow up on https://github.com/pytorch/vision/pull/7592 that makes this Dynamo exportable. Signed-off-by: Edward Z. Yang --- torchvision/models/detection/transform.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchvision/models/detection/transform.py b/torchvision/models/detection/transform.py index 658c9e83455..9c569b0aafb 100644 --- a/torchvision/models/detection/transform.py +++ b/torchvision/models/detection/transform.py @@ -31,8 +31,10 @@ def _resize_image_and_masks( ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]: if torchvision._is_tracing(): im_shape = _get_shape_onnx(image) - else: + elif torch.jit.is_scripting(): im_shape = torch.tensor(image.shape[-2:]) + else: + im_shape = image.shape[-2:] size: Optional[List[int]] = None scale_factor: Optional[float] = None