Added ability to change dataset to MDNTrainer