How to load model saved in local (Visual Transformer ViT)?

Hi there, I have trained a ViT and Fine-Tuned it on Stanford dog dataset.
I saved it in local and now I want to load it (it is a folder with elements in it, as it should be!).
Can you please tell me how to load it? I read online something like:

model = ViTForImageClassification.from_pretrained(‘/drive/MyDrive/vit-stanford-dog-dataset’,local_files_only=True)

but it is not working!


model_directory = “drive/MyDrive/vit-stanford-dog-dataset”
tokenizer = ViTFeatureExtractor.from_pretrained(model_directory)
model = ViTForImageClassification.from_pretrained(model_directory, return_dict=False)

1 Like