As official doc gives example to upscale the image
pixel_shuffle = nn.PixelShuffle(3)
input = torch.randn(1, 9, 4, 4)
output = pixel_shuffle(input)
# torch.Size([1, 1, 12, 12])

As it mutilplies by input, I wanted to use 2 or any other input , giving eror

RuntimeError: invalid argument 2: size ‘[1 x 2 x 2 x 2 x 4 x 4]’ is invalid for input with 144 elements at /opt/conda/conda-bld/pytorch_1524584710464/work/aten/src/TH/THStorage.c:41

I was using on torch size of
torch.randn(1, 512, 6,6)
I wanted to upsale it it (1, 3, 128, 128) by only using pixelShuffle,
But when i use upscaling except 2 all other fails.
At last of upscaling it generates
torch.Size([50, 2, 96, 96])
Help how can i use PixelShuffle to generate image of shape nc=3 and desired height and width.
Thank you.