I have the same requirement but based on yours customised dataset couldn’t figure out how can I print the image name to be able to make sure I am getting the probability class of related image. do you have any idea how can I print the image name ?
this is snippet for the dataset.
        
class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths):   # initial logic happens like transform

        self.image_paths = image_paths
        self.target_paths = target_paths
        self.transforms = transforms.ToTensor()
        self.mapping = {
            0: 0,
            255: 1
     
        }
    def mask_to_class(self, mask):
        for k in self.mapping:
            mask[mask==k] = self.mapping[k]
        return mask
    
    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image = image.convert('L')
        t_image = self.transforms(t_image)
        mask = torch.from_numpy(np.array(mask, dtype=np.uint8)) # this is for my dataset(lv)
        mask = self.mask_to_class(mask)
        mask = mask.long()
        return t_image, mask

    def __len__(self):  # return count of sample we have

        return len(self.image_paths)



folder_data = glob.glob("D:\\Neda\\Pytorch\\U-net\\my_data\\imagesResized\\*.png")
folder_mask = glob.glob("D:\\Neda\\Pytorch\\U-net\\my_data\\labelsResized\\*.png")

# split these path using a certain percentage
len_data = len(folder_data)
print("count of dataset: ", len_data)
train_size = 0.6

train_image_paths = folder_data[:int(len_data*train_size)]
print("count of train images is: ", len(train_image_paths)) # output is 55 image for train
test_image_paths = folder_data[int(len_data*train_size):]
print("count of test images is: ", len(test_image_paths)) # output is 37 image for test

train_mask_paths = folder_mask[:int(len_data*train_size)]
test_mask_paths = folder_mask[int(len_data*train_size):]


train_dataset = CustomDataset(train_image_paths, train_mask_paths)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=2)

test_dataset = CustomDataset(test_image_paths, test_mask_paths)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)


and this is the snippet when I test the test_loader

def test():
    
    model_load = torch.load('model.pth')
    
         #test model
    model_load.eval()
    total = 0
    test_loss = 0
    correct = 0
    count = 0
    #iterate through test dataset
    for ii, data in enumerate(test_loader):
                
                t_image, mask = data
                #print(t_image.shape) # torch.Size([1, 1, 240, 320])
                t_image, mask = t_image.to(device), mask.to(device)
                with torch.no_grad():

                    outputs = model_load(t_image)
                    #print(outputs.shape) # torch.Size([1, 2, 240, 320])
                    test_loss += criterion(outputs, mask).item() / len(test_loader)

                    probs = torch.exp(outputs)
                    probs_numpy = probs.detach().cpu().numpy()
                    b = probs_numpy.tolist()
                    file_path = "/Neda/Pytorch/U-net/output.json" ## your path variable
                    json.dump(b, codecs.open(file_path, 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=4) ### this saves the array in .json format


                    _, predicted = torch.max(outputs.data, 1)  
                    
                    total += mask.nelement()
                    correct += predicted.eq(mask.data).sum().item()
                    accuracy = 100 * correct / total
                    
                    count +=1                
                    print(count, "Test Loss: {:.3f}".format(test_loss), "Test Accuracy: %d %%" % (accuracy))

I can print out image name now by manipulating the print in test function
print(count, "Test Loss: {:.3f}".format(test_loss), "Test Accuracy: %d %%" % (accuracy), os.path.basename(os.path.normpath(test_image_paths[ii])))