def METHOD_NAME(self):
    train_batch_sampler = SamplerHelper(self.train_ds)
    batch_size = 3
    batch_sampler = train_batch_sampler.batch(batch_size)
    for i, sample in enumerate(batch_sampler):
        for j, minibatch in enumerate(sample):
            self.check_output_equal(i * batch_size + j, minibatch)