I have trained a Unet for segmentation with my own dataset, but when I do inference for test images, the predicted masks are the same for some of the input images (which are different).
Here is my predict.py. Is there any problem with my code or any other reasons?
import torch
import sys
from UNet import UNet
import numpy as np
from torchvision import transforms
import PIL.Image as Image
import torch.nn as nn
from torch.autograd import Variable


def main():
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        device = torch.device('cuda:1')
    else:
        device = torch.device('cpu')
    model = UNet().to(device)
    state = torch.load(train_model_path)
    model.load_state_dict(state['state_dict'])

    model.eval()

    test_image = Image.open(predict_image_path)
    test_image = transform(test_image)
    test_image = test_image.unsqueeze(0)
    test_image = Variable(test_image.to(device))

    with torch.no_grad():
        test_scoremap = model(test_image)
        m = nn.Softmax(dim=1)
        test_scoremap = m(test_scoremap)
        test_scoremap = test_scoremap.squeeze(0)

        test_scoremap = test_scoremap.data.cpu().numpy()
        test_pred = test_scoremap[1, :, :] > out_thresh

        test_pred = (test_pred * 255).astype(np.uint8)

        test_img = Image.fromarray(test_pred)
        test_img.save(predict_image_path.split('.')[0] + '_predict.png')


if __name__ == '__main__':
    train_model_path = sys.argv[1]          # the path of trained model file
    predict_image_path = sys.argv[2]        # the path of test image
    out_thresh = 0.5
    transform = transforms.Compose([
            transforms.ToTensor()
        ])
    main()


Here is my code for saving trained model:
checkpoint_path = 'TrainedModels/TrainedModel_{}.pth'.format(datetime.now().strftime('%Y_%m_%d_%H_%M'))
save_checkpoint({
        'epoch': epoch + 1,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'lr_scheduler': scheduler.state_dict(),
        }, checkpoint_path)