Non-Square Images - Githubissues

I'm using this unet architecture to create images, but I'm curious about using non-square images. Current the only problems seems to be in the output_blocks, specifically with th.cat([h, hs.pop(), dim=1]) With non-square images the dimensions aren't correct, which is understandable.

Is there a good way to go about getting this setup to use non-square images? What's interesting is inside the u-net init it takes in an image size but doesn't use it anywhere else