how to train a vision transformer deep learning model in tensorflow