Hi guys!
I’m trying to implement and train the original U-Net model, but I’m stuck in when I’m trying to train the model using the ISBI Challenge Dataset.
According with the original U-Net model, the network outputs an image with 2 channels and size of 388 x 388. So, my data loader for training generates a tensor with size of [batch, channels=1, width=572, height=572] for the input images and [batch, channels=2, width=388, width=388] for target/output images.
My problem actually is that when I’m trying to use the nn.CrossEntropyLoss() the following error is raised:

RuntimeError: invalid argument 3: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THNN/generic/SpatialClassNLLCriterion.c:59

I’m just starting with PyTorch (newbie here)… so, I’ll really appreciate if someone could help me to overcome this problem.
The sourcecode is available on GitHub: https://github.com/dalifreire/cnn_unet_pytorch 7
https://github.com/dalifreire/cnn_unet_pytorch/blob/master/unet_pytorch.ipynb 5
Best regards!