add the train and test loop to train the local vision transformer. create a main.py file which can be called via the cli.