How to use the same batch to train the model?

Luis solution is fine,

For even more customization, consider looking up on how to write a custom Sampler.