Skip to content

Commit

Permalink
Fix potential overflow in convert_image_dtype (pytorch#3107)
Browse files Browse the repository at this point in the history
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'
  • Loading branch information
fmassa authored and vfdev-5 committed Dec 4, 2020
1 parent cb8dcb9 commit f1cd26d
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion torchvision/transforms/functional_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,14 +105,15 @@ 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
if torch.tensor(0, dtype=dtype).is_floating_point():
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
Expand Down

0 comments on commit f1cd26d

Please sign in to comment.