imjeffhi
/

pokemon_classifier

Image Classification

Inference Endpoints

Model card Files Files and versions Community

imjeffhi commited on Jan 1, 2022

Commit

83deabd

•

1 Parent(s): f645892

Update README.md

Files changed (1) hide show

README.md +2 -2

README.md CHANGED Viewed

@@ -15,8 +15,8 @@ import torch
 # Loading in Model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = ViTForImageClassification.from_pretrained( "./PokemonModel").to(device)
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
 # Caling the model on a test image
 img = Image.open('test.jpg')

 # Loading in Model
 device = "cuda" if torch.cuda.is_available() else "cpu"
+model = ViTForImageClassification.from_pretrained( "imjeffhi/pokemon_classifier").to(device)
+feature_extractor = ViTFeatureExtractor.from_pretrained('imjeffhi/pokemon_classifier')
 # Caling the model on a test image
 img = Image.open('test.jpg')