Convert pyTorch to CoreML

I converted this sample code to CoreML using coremltools 7 and the output is showing as MultiArray (Float16 1 × 256 × 128 × 128). I am after a MultiArray of 512 x 512 similar to the input Image (Color 512 × 512).

Do I need to add layers to do this ?

https://pytorch.org/tutorials/advanced/neural_style_tutorial.html

The model definition shows as:

model=Sequential(
(0): Normalization()
(conv_1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(style_loss_1): StyleLoss()
(relu_1): ReLU()
(conv_2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(style_loss_2): StyleLoss()
(relu_2): ReLU()
(pool_2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv_3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(style_loss_3): StyleLoss()
(relu_3): ReLU()
(conv_4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(content_loss_4): ContentLoss()
(style_loss_4): StyleLoss()
(relu_4): ReLU()
(pool_4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv_5): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(style_loss_5): StyleLoss()
)