Conditional cast for onnx.

Author: datumbox

torchvision/models/detection/transform.py

@@ -11,27 +11,28 @@
 
 @torch.jit.unused
 def _get_shape_onnx(image):
-    # type: (Tensor) -> Tensor
     from torch.onnx import operators
     return operators.shape_as_tensor(image)[-2:]
 
+@torch.jit.unused
+def _float_to_tensor_onnx(v):
+    return torch.tensor(v)
+
 
 def _resize_image_and_masks(image, self_min_size, self_max_size, target):
     # type: (Tensor, float, float, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]
 
     if torchvision._is_tracing():
         im_shape = _get_shape_onnx(image)
-        cast_to_float = False
     else:
         im_shape = torch.tensor(image.shape[-2:])
-        cast_to_float = True
 
     min_size = torch.min(im_shape).to(dtype=torch.float32)
     max_size = torch.max(im_shape).to(dtype=torch.float32)
-    scale_factor = torch.min(self_min_size / min_size, self_max_size / max_size)
+    scale_factor = float(torch.min(self_min_size / min_size, self_max_size / max_size))
 
-    if cast_to_float:
-        scale_factor = float(scale_factor)
+    if torchvision._is_tracing():
+        scale_factor = _float_to_tensor_onnx(scale_factor)
 
     image = torch.nn.functional.interpolate(
         image[None], scale_factor=scale_factor, mode='bilinear', recompute_scale_factor=True,