Running code on custom segmentation dataset

[SinghTejwinder]

Hi, I am trying to run this code on my custom image segmentation dataset which has 13 label classes but I am getting the following error :

RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size: : [4, 768, 768, 3]

I checked for the output and target tensor sizes and they are

-------------Output------------ torch.Size([4, 13, 768, 768])
-------------Target------------ torch.Size([4, 768, 768, 3])

I am not getting why the code is predicting an image with 13 channels. Can you please let me know what changes do I need to make in order to run this code.