Could you explain your input data a bit?
I.e. what kind of input images and masks are you using?
Somewhere in your code you are most likely using a Dataset, which loads the image and the corresponding mask. Are you assuming the mask size is different than the image size?
If so, would resizing work?