That makes a ton of sense - after researching a bit, I think I will go with the approach where I create a separate Dataset for each input class, wrapped in a ConcatDataset. Then the DataLoader would call on a custom sampler which would alternate between pulling from the different sub-Datasets…
Thanks @ptrblck!