This seems be be along the lines of the question I asked here: Batch grouping - fastai users - Deep Learning Course Forums. Basically, in my training set I have some (slightly) different training sets and only sets of the same kind should be combined in a batch (in my case due to the dimensions of the data varying slightly). However, I can’t quite seem to connect the dots and see if I can make the approach you (@muellerzr) outline work.