Add the dataset and the training script

Add the model and scheduler