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.

image417×693 14.1 KB