Learning how to train models on given dataset