Fix batch shuffle test
1.0.1