Hello dear engineers,
Hope you are all doing well. I am still in my first steps in both pytorch and programming. My task is a binary segmentation problem which I would like to approach it as a binary classification. My data are RGB of size (360, 360, 3) and the labels are black and white. I have adopted the following simple CNN model for the training process.
class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()

        self.cnn1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=0)
        self.relu1 = nn.ReLU()
        self.maxpool1 = nn.MaxPool2d(kernel_size=2)

        self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=0)
        self.relu2 = nn.ReLU()
        self.maxpool2 = nn.MaxPool2d(kernel_size=2)

        self.dropout = nn.Dropout(p=0.5)

        #self.fc1 = nn.Linear(32 * 4* 4, 2)
        self.fc1 = nn.Linear(32 * 87 * 87, 2)

    def forward(self, x):
        out = self.cnn1(x)
        out = self.relu1(out)
        out = self.maxpool1(out)
        out = self.cnn2(out)
        out = self.relu2(out)
        out = self.maxpool2(out)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.fc1(out)

        return out

My training function is as follows.

opt = Config()
device = torch.device('cuda') if opt.use_gpu else torch.device('cpu')
num_classes = 2

train_dataset = Dataset(opt.train_input, opt.train_label)
print(len(train_dataset))
# print(train_dataset.__getitem__(0)[0].shape)
train_loader = data.DataLoader(train_dataset,
                                batch_size=opt.train_batch_size,
                                shuffle=True,
                                num_workers=opt.num_workers)

valid_dataset = Dataset(opt.valid_input, opt.valid_label)

print(len(valid_dataset))
valid_loader = data.DataLoader(valid_dataset,
                                batch_size=opt.train_batch_size,
                                shuffle=True,
                                num_workers=opt.num_workers)


CE = nn.CrossEntropyLoss().to(device)


def train(train_loader, model, optimizer, criterion=CE):

    model.train()

    for i, (input, target) in enumerate(tqdm(train_loader)):
        # if idx.size(0) != opt.train_batch_size:
        #     break

        input = torch.autograd.Variable(input.to(device))
        target = torch.autograd.Variable(target.to(device))
        # input = input.to(device)
        # target = target.to(device)
        target = target.long()
        output = model(input)
        #print(target.shape, output.shape)
        #print(torch.max(target, 1))
        loss = criterion(output, target[:,2,0,0])

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0

    for i, (input, target) in enumerate(tqdm(test_loader)):
        input = torch.Tensor(input).to(device)
        target = torch.autograd.Variable(target).to(device)

        total += target.size(0)
        output = model(input)
        print(target.shape, output.shape)
        _, predicted = output.max(1)
        correct += predicted.eq(target[:,2,0,0]).sum().item()

    accuracy = 100. * correct / total

    return accuracy


def main_ce():
    model_ce = CNNModel().to(device)
    best_valid_acc = 0

    for epoch in range(20):
        print("epoch=", epoch)
        #print("r=", r)
        learning_rate = 1e-4
        optimizer_ce = torch.optim.Adam(model_ce.parameters(), lr=learning_rate)
        print("traning model_ce...")
        train(train_loader=train_loader, model=model_ce, optimizer=optimizer_ce)
        print("validating model_ce...")
        valid_acc = test(model=model_ce, test_loader=valid_loader)
        print('valid_acc=', valid_acc)
        if valid_acc >= best_valid_acc:
            best_valid_acc = valid_acc
            torch.save(model_ce, './model_ce_' + str(1) + '_' + str(1) + '_' + str(1))
            print("saved.")
        print("testing model_ce...")
        test_acc = test(model=model_ce, test_loader=valid_loader)
        print('test_acc=', test_acc)



if __name__ == '__main__':
    main_ce()

The code can run without bugs. However, the accuracy remains unchanged after training for 20 epochs. I feel like there is something wrong with the codes. Please, could anyone help me check it out? Any comments or suggestions would be highly appreciated.