The train.py script says that currently only batch size of 1 is implemented. Is arbitrary batch size implemented now?