Thank you for your reply.
Sorry for incomplete information
This is a semantic segmentation problem
Input data: Image is in .jpg format. I use transform.ToTensor() to convert this.
Target: Annotated .png image with 12 classes. I convert it using “torch.from_numpy(np.array(mask)).long()”
Following is the code I am working on.
data loader class
class GetDataset(torch.utils.data.Dataset):
def init(self, root, transformation=None):
self.root = root
self.transformation = transformation
# load all image files, sorting them to
# ensure that they are aligned
self.imgs = list(sorted(os.listdir(os.path.join(root, “JPEGImages_test”))))
self.masks = list(sorted(os.listdir(os.path.join(root, “SegmentationClass_test”))))
def __getitem__(self, idx):
    # load images ad masks
    img_path = os.path.join(self.root, "JPEGImages_test", self.imgs[idx])
    mask_path = os.path.join(self.root, "SegmentationClass_test", self.masks[idx])
    print(idx)
    #import pdb
    #pdb.set_trace()
    img = Image.open(img_path)
    
    mask = Image.open(mask_path)
    
    img_t = self.transformation(img)
    mask_t = torch.from_numpy(np.array(mask)).long()
    return img_t, mask_t 

def __len__(self):
    return len(self.imgs)  

==> Tried both losses
#criterion = nn.CrossEntropyLoss()
criterion = nn.NLLLoss()
Train function
def train_model(model, criterion, optimizer, dataloaders, scheduler, num_epochs=1):
import copy
best_model_wts = copy.deepcopy(model.state_dict())
print(num_epochs)
for epoch in range(num_epochs):
    print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    print('-' * 10)
    model.train()  # Set model to training mode

    for inputs, labels in dataloaders:
        print("reached")
        
        inputs = inputs.to(device)
        type(input)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        #import pdb
        #pdb.set_trace()
        loss = criterion(outputs['out'], labels)
        loss.backward()
        optimizer.step()

    best_model_wts = copy.deepcopy(model.state_dict())

# load best model weights
model.load_state_dict(best_model_wts)
return model

#Train model
model_ft = train_model(model_segdl, criterion, optimizer_ft,data_loader, exp_lr_scheduler, num_epochs=1)
model_ft.eval()
x = torch.rand(1, 3, 300, 400)
predictions = model_ft(x)
predictions[‘out’].size()
Out[59]: torch.Size([1, 12, 300, 400])
Is this correct ? If not what I am missing