Bugfix: train now batches properly