Fix batch sampler (#2097)

Commit

2 years ago

Fix batch sampler (#2097) * Fix batch sampler * Clean * Fix tests * Fix * Better comment * Base case