Hi,
what resolution do your images have? Right now you are creating an UNet with depth=10. That means according to the model that



 copperwiring:

input size needs to be (depth-1) times divisible by 2



which equals 512. The original UNet has a depth of 4 if I’ve counted correctly. Maybe decreasing the depth of your model will solve the problem, but that’s only a wild guess.
EDIT:
You are using BCEWithLogitsLoss which seems to be okay according to this post 4. But your UNet-model says in the comments, that
        # No softmax is used. This means we need to use
        # nn.CrossEntropyLoss is your training script,
        # as this module includes a softmax already.

Could this be a problem?