Add support for mapping datasets for training
Add support for mapping datasets for training