Mismatch output size on last layer

I am using UNet in hub (U-Net for brain MRI | PyTorch) for the depth estimation task. I have no clue what is going wrong with this error

RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of size: : [64, 1, 256, 256]

Some context on my case: I am using pytorch lightning as a trainer and keeping the hyper params of model as suggested on the link above.

Could you describe your use case, which line of code raises the error, and what the target represents?

Here is the full error log. The target is a depth image with values in the range [-1, 1] in size 256x256.

The shape mismatch is raised in nn.NLLLoss so check the docs and make sure the model output and target have the expected shapes.
I would guess you might need to squeeze dim1 in the target but since you didn’t explain the use case it’s just my guess.

For my use case, I am trying to use UNet for the depth estimation task. And now I understand, my label is an image size 256x256, and my model is only outputting a single label. How could I change the last layer of UNet to output a result that matches my label?

The problem is that I apply a training procedure for classification in my regression task. I fix the loss function and everything is ok now.

Ah OK, I guess you are using nn.MSELoss now and it works fine?

1 Like

You guess it right. I forgot to check the loss function of the training pipeline at first.