Move the asserts in shape functions upsample_nearest_2d op. (#85801)

The assert check are moved to top and the function now returns out. This is needed by the downstream torch-mlir project to correctly determine the output type.

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/85801
Approved by: https://github.com/eellison
This commit is contained in:
Prashant Kumar
2022-09-30 18:30:06 +00:00
committed by PyTorch MergeBot
parent b60ad2e529
commit 7ddf167ba5

View File

@ -355,6 +355,10 @@ def upsample_nearest2d(
out: List[int] = []
out.append(input[0])
out.append(input[1])
if (scale_factors is None and output_size is None):
assert 0, "Either output_size or scale_factors must be presented"
if output_size is not None:
assert (
scale_factors is None
@ -362,7 +366,6 @@ def upsample_nearest2d(
assert len(output_size) == 2
out.append(output_size[0])
out.append(output_size[1])
return out
if scale_factors is not None:
assert (
@ -371,8 +374,8 @@ def upsample_nearest2d(
assert len(scale_factors) == 2
out.append(int(input[2] * scale_factors[0]))
out.append(int(input[3] * scale_factors[1]))
return out
assert 0, "Either output_size or scale_factors must be presented"
return out
def mm(self: List[int], mat2: List[int]):