Updates vit_jax_augreg.ipynb
Updates vit_jax_augreg.ipynb

Vastly improved finetuning section (incl. TensorBoard, explanation of
config values and fine-tuning from local file system / Google Drive).

google-research/vision_transformer