Hugging Face's logo

imjeffhi
/
pokemon_classifier

Image Classification
Transformers
PyTorch
vit
Model card Files
xet
 Community
3
imjeffhi commited on
Commit
83deabd
·
1 Parent(s): f645892

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -2
README.md CHANGED
@@ -15,8 +15,8 @@ import torch
15
 
16
  # Loading in Model
17
  device = "cuda" if torch.cuda.is_available() else "cpu"
18
- model = ViTForImageClassification.from_pretrained( "./PokemonModel").to(device)
19
- feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
20
 
21
  # Caling the model on a test image
22
  img = Image.open('test.jpg')
 
15
 
16
  # Loading in Model
17
  device = "cuda" if torch.cuda.is_available() else "cpu"
18
+ model = ViTForImageClassification.from_pretrained( "imjeffhi/pokemon_classifier").to(device)
19
+ feature_extractor = ViTFeatureExtractor.from_pretrained('imjeffhi/pokemon_classifier')
20
 
21
  # Caling the model on a test image
22
  img = Image.open('test.jpg')