Avoid creating a tensor of shape when not tracing

This is a small follow up on #7592 that makes this Dynamo exportable. Signed-off-by: Edward Z. Yang <ezyang@meta.com>
pytorch · Sep 7, 2023 · d50177f · d50177f
1 parent e13b8f5
commit d50177f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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