Conditional cast.

Author: datumbox

torchvision/models/detection/transform.py

@@ -21,12 +21,17 @@ def _resize_image_and_masks(image, self_min_size, self_max_size, target):
 
     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 = float(torch.min(self_min_size / min_size, self_max_size / max_size))
+    scale_factor = torch.min(self_min_size / min_size, self_max_size / max_size)
+
+    if cast_to_float:
+        scale_factor = float(scale_factor)
 
     image = torch.nn.functional.interpolate(
         image[None], scale_factor=scale_factor, mode='bilinear', recompute_scale_factor=True,