Fix `_upsample_2d` (#535)

* Fix _upsample_2d

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar 2022-09-19 15:52:52 +02:00 committed by GitHub
parent b8d1f2d344
commit b17d49f863
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 2 deletions

View File

@ -149,7 +149,6 @@ class FirUpsample2D(nn.Module):
stride = (factor, factor)
# Determine data dimensions.
stride = [1, 1, factor, factor]
output_shape = ((x.shape[2] - 1) * factor + convH, (x.shape[3] - 1) * factor + convW)
output_padding = (
output_shape[0] - (x.shape[2] - 1) * stride[0] - convH,
@ -161,7 +160,7 @@ class FirUpsample2D(nn.Module):
# Transpose weights.
weight = torch.reshape(weight, (num_groups, -1, inC, convH, convW))
weight = weight[..., ::-1, ::-1].permute(0, 2, 1, 3, 4)
weight = torch.flip(weight, dims=[3, 4]).permute(0, 2, 1, 3, 4)
weight = torch.reshape(weight, (num_groups * inC, -1, convH, convW))
x = F.conv_transpose2d(x, weight, stride=stride, output_padding=output_padding, padding=0)