How to implement custom batch sizes

data
1

I am trying to do a Listwise LTR NN, where I am predicting a rank based on the QID.

I currently have this custom datset but now need to make sure each batch is only relevant to the QID and doesn’t have other QID in same batch

image

Any pointers would be great please