RuntimeError: Function AddmmBackward returned an invalid gradient at index 1 - got [16, 2048] but expected shape compatible with [16, 32768]

Hello. I am trying to train a plain VAE on my custom dataset. I’m transforming the images into a (3,256,256) tensor and feeding it into the dataloader. The input and output tensors match, but when I run loss.backward(), it gives me the below error message:

“RuntimeError: Function AddmmBackward returned an invalid gradient at index 1 - got [16, 2048] but expected shape compatible with [16, 32768]”

Can you please clarify the issue and propose a solution? Thank you very much in advance.

Below are additional snips of the code…

Please don’t post images as we cannot copy/paste the code in order to debug and the forum search also won’t be able to index your post.
You can post code snippets by wrapping them into three backticks ```, which would make debugging easier :wink:
Make sure the code is executable (i.e. we can copy it and directly execute).