jacoballessio
/

ai-image-detect

Image Classification

Model card Files Files and versions

jacoballessio commited on Apr 7, 2024

Commit

5204502

·

verified ·

1 Parent(s): b9fedf9

Update README.md

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -11,7 +11,7 @@ from torchvision import transforms
 from transformers import ViTForImageClassification, ViTImageProcessor
 # Load the trained model
-model_path = 'trained_modelBEST.pth'
 model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
 model.classifier = torch.nn.Linear(model.classifier.in_features, 2)
 model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

 from transformers import ViTForImageClassification, ViTImageProcessor
 # Load the trained model
+model_path = 'vit_model.pth'
 model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
 model.classifier = torch.nn.Linear(model.classifier.in_features, 2)
 model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))