Visually aligning 2 sets of data

Let’s say that I have 2 sets of data (images). Domain A and domain B. I randomly sample a batch_size of samples from domain A. I want to find batch_size of samples in domain B, which are the most similar (by some metric) to batch_size samples in domain A. The most similar samples in domain B are sampled from a variable-sized random subset (no to do an extensive search of the whole sample space).

Basically, if I sample 16 images in domain A, I want to find the most similar example for each sample in domain B (from a random subset e.g. 100 images). 16 images from A and 16 (closest pair-wise) then present input to the model.

What would be the cleanest approach to implement such a pipeline in PyTorch (using Datasets, DataLoaders, Samplers, etc.)