[Flax] ViT training example
# What does this PR do?

This PR adds image classification script for fine-tuning flax vit.

For faster processing and loading , torch dataloader is used, since this can become a bottleneck on TPU.
Can you run `make style` ?