'''
TODO: Add the training script here
'''