Looks like you crash if width and height isn't equal on custom datasets. I worked around by forcing my 640x360 images to 640x640.