AutoEncoderKL output tensor dimension mismatch with Input

[shankartmv]

I am trying to train a AutoEncoderKL model on RGB images with the following dimensions (3,1225,966). Here is the code that I use ( similar to what's there in tutorials/generative/2d_ldm/2d_ldm_tutorial.ipynb ).
autoencoderkl = AutoencoderKL(
spatial_dims=2,
in_channels=3,
out_channels=3,
num_channels=(128, 256, 384),
latent_channels=8,
num_res_blocks=1,
attention_levels=(False, False, False),
with_encoder_nonlocal_attn=False,
with_decoder_nonlocal_attn=False,
)
autoencoderkl = autoencoderkl.to(device)

Error is reported at line 27 (Train Model - as in the tutorials notebook)
recons_loss = F.l1_loss(reconstruction.float(), images.float()) RuntimeError: The size of tensor a (964) must match the size of tensor b (966) at non-singleton dimension 3

Using pytorchinfo package , I was able to print the model summary and can find the discrepancy in the upsampling layer.

===================================================================================================================
Layer (type:depth-idx) Input Shape Output Shape Param #
===================================================================================================================
AutoencoderKL [1, 3, 1225, 966] [1, 3, 1224, 964] --
├─Encoder: 1-1 [1, 3, 1225, 966] [1, 8, 306, 241] --
│ └─ModuleList: 2-1 -- -- --
│ │ └─Convolution: 3-1 [1, 3, 1225, 966] [1, 128, 1225, 966] 3,584
│ │ └─ResBlock: 3-2 [1, 128, 1225, 966] [1, 128, 1225, 966] 295,680
│ │ └─Downsample: 3-3 [1, 128, 1225, 966] [1, 128, 612, 483] 147,584
│ │ └─ResBlock: 3-4 [1, 128, 612, 483] [1, 256, 612, 483] 919,040
│ │ └─Downsample: 3-5 [1, 256, 612, 483] [1, 256, 306, 241] 590,080
│ │ └─ResBlock: 3-6 [1, 256, 306, 241] [1, 384, 306, 241] 2,312,576
│ │ └─GroupNorm: 3-7 [1, 384, 306, 241] [1, 384, 306, 241] 768
│ │ └─Convolution: 3-8 [1, 384, 306, 241] [1, 8, 306, 241] 27,656
├─Convolution: 1-2 [1, 8, 306, 241] [1, 8, 306, 241] --
│ └─Conv2d: 2-2 [1, 8, 306, 241] [1, 8, 306, 241] 72
├─Convolution: 1-3 [1, 8, 306, 241] [1, 8, 306, 241] --
│ └─Conv2d: 2-3 [1, 8, 306, 241] [1, 8, 306, 241] 72
├─Convolution: 1-4 [1, 8, 306, 241] [1, 8, 306, 241] --
│ └─Conv2d: 2-4 [1, 8, 306, 241] [1, 8, 306, 241] 72
├─Decoder: 1-5 [1, 8, 306, 241] [1, 3, 1224, 964] --
│ └─ModuleList: 2-5 -- -- --
│ │ └─Convolution: 3-9 [1, 8, 306, 241] [1, 384, 306, 241] 28,032
│ │ └─ResBlock: 3-10 [1, 384, 306, 241] [1, 384, 306, 241] 2,656,512
│ │ └─Upsample: 3-11 [1, 384, 306, 241] [1, 384, 612, 482] 1,327,488
│ │ └─ResBlock: 3-12 [1, 384, 612, 482] [1, 256, 612, 482] 1,574,912
│ │ └─Upsample: 3-13 [1, 256, 612, 482] [1, 256, 1224, 964] 590,080
│ │ └─ResBlock: 3-14 [1, 256, 1224, 964] [1, 128, 1224, 964] 476,288
│ │ └─GroupNorm: 3-15 [1, 128, 1224, 964] [1, 128, 1224, 964] 256
│ │ └─Convolution: 3-16 [1, 128, 1224, 964] [1, 3, 1224, 964] 3,459
===================================================================================================================
Total params: 10,954,211
Trainable params: 10,954,211
Non-trainable params: 0
Total mult-adds (Units.TERABYTES): 3.20
===================================================================================================================
Input size (MB): 14.20
Forward/backward pass size (MB): 26803.57
Params size (MB): 43.82
Estimated Total Size (MB): 26861.59
===================================================================================================================

[shankartmv]

After some debugging I figured out a way to get around this problem. By resizing my images to standard 3:2 aspect ratio, (1024*720) I can see that the input and output shapes (obtained from pytorch.summary) of my AutoEncoderKL is consistent. But anyways, I would like to know the reason behind this error.

[minghuix98]

I believe this is caused by downsampling and upsampling on data with a nan 2 power dimension.

[virginiafdez]

I think this happens cause you have downsamplings that divide the spatial dimensions by 2 and upsample, so unless you play around with the paddings and strides to make sure things end up having the same size, you might run into errors. I would recommend simply padding your inputs to a size that is consistently divisible by 2.