I think that’s what I’ll fall back to, thanks. I was hoping there would be a more efficient way to do this using DataLoader - assigning individual samples to custom batches based on some criteria feels like it should not be such a unique use case.