google
/
vit-base-patch16-224
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update README

This commit is contained in:
Niels Rogge 2021-03-30 18:14:02 +00:00 committed by huggingface-web
parent e457400e6d
commit a46a961b1e
1 changed files with 2 additions and 1 deletions

3
README.md
View File

@ -45,7 +45,8 @@ inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
predicted_class = logits.argmax(-1).item()
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])
```
Currently, both the feature extractor and model support PyTorch. Tensorflow and JAX/FLAX are coming soon, and the API of ViTFeatureExtractor might change.