You could sample some variable from [0, len(val_loader)-1] and compare it directly to i_val inside the data loader loop.
Then just use the same variable to get the corresponding image.
Alternatively, you could also remove the plotting completely from the data loader loop and just sample from the Dataset directly using a random number.