@ptrblck Thank you for your help in working with Cross Entropy Loss. However, the predictions (for three classes) seem to be too off.
When I use the code below:
            print("shape of model input images:", images.shape)
            print("shape of masks_input:", masks_input.shape)
            print('masks input:')
            print(masks_input)
            outputs = model(images)  #model doesn't have softmax and images are normalized
            masks = torch.argmax(masks_input, dim=1)
            loss = loss_new(outputs, masks) 
            preds = torch.argmax(outputs, dim=1)
            print("shape of preds:", preds.shape)
            print('preds output')
            print(preds)

I get the following pred output. Basically, 95% of it is being predicted of Class 1 -
shape of model input images: torch.Size([1, 5, 256, 256])
shape of masks_input: torch.Size([1, 3, 256, 256])
masks input:
tensor([[[[0., 0., 0.,  ..., 1., 1., 1.],
          [0., 0., 0.,  ..., 1., 1., 1.],
          [0., 0., 0.,  ..., 1., 1., 0.],
          ...,
          [0., 0., 1.,  ..., 0., 0., 0.],
          [0., 1., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]],

         [[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 1.],
          ...,
          [1., 1., 0.,  ..., 0., 0., 1.],
          [1., 0., 1.,  ..., 0., 0., 0.],
          [1., 1., 0.,  ..., 0., 0., 0.]],

         [[1., 1., 1.,  ..., 0., 0., 0.],
          [1., 1., 1.,  ..., 0., 0., 0.],
          [1., 1., 1.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 1., 1., 0.],
          [0., 0., 0.,  ..., 1., 1., 1.],
          [0., 0., 1.,  ..., 1., 1., 1.]]]])
shape of preds: torch.Size([1, 256, 256])
preds output
tensor([[[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 2],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 2],
         [0, 0, 0,  ..., 0, 0, 0]]])