Hi, Thanks for responding.

Well I tried what you suggested but didn’t work. But it seems that using transform.PILToTensor (for masks) instead of transforms.Tensor works fine as the masks don’t lose information then. Can you tell me why ?
Another thing I am working on CamVid dataset with 32 classes. Initial mask provided was of size [h,w,3] I converted it into [h,w,32] where each channel represent corresponding class. Is this approach correct ?
Also, I am not sure about if I should use Sigmoid or Softmax function. And Currently I am using loss BCE
The prediction on test is not good



UNet:
def Double_Conv(input_channel, output_channel):
     return nn.Sequential(
          nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1),
          nn.BatchNorm2d(output_channel),
          nn.ReLU(inplace=True),
          nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1),
          nn.BatchNorm2d(output_channel),
          nn.ReLU(inplace=True))

class UNet(nn.Module):
      def __init__(self, no_of_cls):
        super().__init__()

        self.down1 = Double_Conv(3, 64)
        self.down2 = Double_Conv(64, 128)
        self.down3 = Double_Conv(128, 256)
        self.down4 = Double_Conv(256, 512)
        self.down5 = Double_Conv(512, 1024)

        self.max_pool = nn.MaxPool2d(2)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.up4 = Double_Conv(512 + 1024, 512)
        self.up3 = Double_Conv(256 + 512, 256)
        self.up2 = Double_Conv(128 + 256, 128)
        self.up1 = Double_Conv(64 + 128, 64)

        self.lastConv = nn.Conv2d(64, no_of_cls, 1)

def forward(self, x):
  conv1 = self.down1(x)
  x = self.max_pool(conv1)

  conv2 = self.down2(x)
  x = self.max_pool(conv2)

  conv3 = self.down3(x)
  x = self.max_pool(conv3)

  conv4 = self.down4(x)
  x = self.max_pool(conv4)

  x = self.down5(x)

  x = self.upsample(x)
  x = torch.cat([x, conv4], dim=1)

  x = self.up4(x)
  x = self.upsample(x)
  x = torch.cat([x, conv3], dim=1)

  x = self.up3(x)
  x = self.upsample(x)
  x = torch.cat([x, conv2], dim=1)

  x = self.up2(x)
  x = self.upsample(x)
  x = torch.cat([x, conv1], dim=1)

  x = self.up1(x)
  out = self.lastConv(x)
  return torch.sigmoid(out)