batch_size for MNISTDataModule

[dedeswim]

🚀 Feature

Add the batch_size parameter to the MNISTDataModule and BinaryMNISTDataModule.

Motivation

When using the MNISTDataModule there is no way to set the batch size if it is used directly within PyTorch Lightning (i.e. as argument to Trainer.fit or as _datamodule field inside a LightningModule).

Pitch

I would like to be able to set the batch size when initializing an MNISTDataModule as I can with many other DataModules right now (like CIFAR10DataModule or ImagenetDataModule).

Alternatives

An alternative to set the batch size would be not to feed the DataModule to the trainer directly (or using it as _datamodule field), but to use the separate test_dataloader, train_dataloader and val_dataloader methods separately. However I think ti would be against one of the points of using a DataModule.

Additional context

A possible implementation could be like that shown on Lightning's docs. I would be available to open a PR and work on this.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[nateraw]

Done 😄 ...you got num_workers for free too!

[hecoding]

Hi, it's not actually working. Still taking batch_size from the method parameter.
I could make a silly PR if anyone interested.

[hecoding]

Well I was the one interested :)
Changes are in the PR above ^

I need some reviewers. Should I create a new issue or something? (not much experience with contributions)