Support customizing batch size in train test (#1505)
Summary:
Also allow users to setup default batch size for train tests.
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1505
Reviewed By: aaronenyeshi
Differential Revision: D45356936
Pulled By: xuzhao9
fbshipit-source-id: da0650c16f09b08866bcb4a3fe80875d8fa81482
