From f1cd26d443171d9adeeb08badd13d714645c1063 Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Thu, 3 Dec 2020 19:40:10 +0100 Subject: [PATCH] Fix potential overflow in convert_image_dtype (#3107) We could have errors such as aten/src/ATen/native/cpu/PowKernel.cpp:41:5: runtime error: 5.7896e+76 is outside the range of representable values of type 'float' --- torchvision/transforms/functional_tensor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchvision/transforms/functional_tensor.py b/torchvision/transforms/functional_tensor.py index 0c72a745bba..ebc178ac561 100644 --- a/torchvision/transforms/functional_tensor.py +++ b/torchvision/transforms/functional_tensor.py @@ -105,7 +105,6 @@ def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) - return result.to(dtype) else: input_max = _max_value(image.dtype) - output_max = _max_value(dtype) # int to float # TODO: replace with dtype.is_floating_point when torchscript supports it @@ -113,6 +112,8 @@ def convert_image_dtype(image: torch.Tensor, dtype: torch.dtype = torch.float) - image = image.to(dtype) return image / input_max + output_max = _max_value(dtype) + # int to int if input_max > output_max: # factor should be forced to int for torch jit script