Same prediction result for different inputs

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)

I dont see anything obviously wrong with what you are doing. Can you try re-save test_image to disk along-side test_img, and see if it has something to do with loading paths?
Also the function body of m i.e. of m(test_scoremap) is not in your snippet, check if there’s any bugs in that.

I found a problem with my csv file for dataset. All the ground truth filenames are pointing to a single mask.