I am trying to predict a target image from an input image and then display both the input and prediction. For some reason, test_images seems to be reshuffled in the process so I can’t get them to line up. test_dataset contains a tuple with two tensors, one of the input images and the other of the target images.

test_dataset_batched = test_dataset.batch(batch_size=1)
test_predictions = model.predict(test_dataset_batched.map(lambda x, y: x))
test_images = test_dataset.map(lambda x, y: x)
test_images = list(test_images.as_numpy_iterator())
fig, ax = plt.subplots(1, 2, figsize=(10, 10))
ax[0].imshow(test_predictions[0], cmap='gray')
ax[1].imshow(test_images[0], cmap='gray')
plt.show()