Unverified Commit 2312b27f authored by Pedro Cuenca's avatar Pedro Cuenca Committed by GitHub
Browse files

Interpolate fix on cuda for large output tensors (#10067)



* Workaround for upscale with large output tensors.

Fixes #10040.

* Fix scale when output_size is given

* Style

---------
Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent 6db33337
...@@ -165,6 +165,14 @@ class Upsample2D(nn.Module): ...@@ -165,6 +165,14 @@ class Upsample2D(nn.Module):
# if `output_size` is passed we force the interpolation output # if `output_size` is passed we force the interpolation output
# size and do not make use of `scale_factor=2` # size and do not make use of `scale_factor=2`
if self.interpolate: if self.interpolate:
# upsample_nearest_nhwc also fails when the number of output elements is large
# https://github.com/pytorch/pytorch/issues/141831
scale_factor = (
2 if output_size is None else max([f / s for f, s in zip(output_size, hidden_states.shape[-2:])])
)
if hidden_states.numel() * scale_factor > pow(2, 31):
hidden_states = hidden_states.contiguous()
if output_size is None: if output_size is None:
hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest") hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment