I know the type of sampling you’re talking about, thanks for the clarification. Since you have a good grasp on what the training loop should look like, I’d recommend just writing a custom one. Flux’s default is very constrained and this is one of those cases where the intent of the balancing might be clearer if the iteration over multiple dataloaders is made explicit.