The train_test_split feature is not currently supported within the load_mnist_dataset. The user will have to call the split before calling the wrapper function.

It will be a better design to include the split within the function such that both the train and test datasets or Dataloaders are retrieved, split, and normalized with one function only. Another option could be to let load_mnist_dataset take an argument that mentions if the user is loading the train or test dataloader. To achieve this, logic inside the wrapper function must be split into two sub-functions or conditionals.

Ideas for better design are welcome.