Do you mean something like this?
class InferenceData(Dataset):
    def __init__(self, transform=None, phase_train=False, data_dir=None):

        self.phase_train = phase_train
        self.transform = transform
        self.data_dir = data_dir
        
        self.img_dir_test = os.path.join(data_dir, 'data', 'inference', 'images')
        self.depth_dir_test = os.path.join(data_dir, 'data', 'inference', 'depths')

    def __len__(self):
        if self.phase_train:
            img_files_train = os.listdir(self.img_dir_train)
            return len(img_files_train)
        else:
            img_files_test = os.listdir(self.img_dir_test)
            return len(img_files_test)

    def __getitem__(self, idx):
        if self.phase_train:
            img_dir = self.img_dir_train
            depth_dir = self.depth_dir_train
        else:
            img_dir = self.img_dir_test
            depth_dir = self.depth_dir_test
            
        img_paths = glob.glob(os.path.join(img_dir, '*.png'))
        depth_paths = glob.glob(os.path.join(depth_dir, '*.png'))
        
        depth = np.asarray(Image.open(depth_paths[idx]))
        image = np.asarray(Image.open(img_paths[idx]))
        image = skimage.transform.resize(image, (image_h, image_w), order=1,
                                         mode='reflect', preserve_range=True)
        depth = skimage.transform.resize(depth, (image_h, image_w), order=0,
                                         mode='reflect', preserve_range=True)

        image = image / 255
        image = torch.from_numpy(image).float()
        depth = torch.from_numpy(depth).float()
        image = image.permute(2, 0, 1)
        depth.unsqueeze_(0)
        image = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image)
        depth = torchvision.transforms.Normalize(mean=[19050], std=[9650])(depth)
        
        sample = {'image': image, 'depth': depth}
        
        if self.transform:
            sample = self.transform(sample)

        return sample, idx, img_paths[idx], depth_paths[idx]

If yes, I am having problems referencing the individual transformed images.
My inference code looks something like this.
for batch_idx, (sample, idx, img_paths, depth_paths) in enumerate(test_loader):
    
        image = sample['image']
        depth = sample['depth']
        
        if batch_idx % 1 == 0:
            print('No. of Batches Done: [{0}/{1}]\t'.format(batch_idx,len(test_loader)))

        i = 0
        
        for im, d in zip(image, depth):
            print(im[i].shape)
            print(d[i].shape)
            im[i] = im[i].to(device).unsqueeze_(0)
            d[i] = d[i].to(device).unsqueeze_(0)
            pred = model(im, d)
            output = utils.color_label(torch.max(pred, 1)[1]+1)[0]