model.py

from transformers import ViTFeatureExtractor, ViTForImageClassification

def create_ViT():

  extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
  model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

  return model