Update README

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

View File

@ -41,7 +41,7 @@ url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
inputs = feature_extractor(images=image)
inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes