facebook
/

vit-mae-base

Model card Files Files and versions

nielsr HF Staff commited on Jan 19, 2022

Commit

54a0cf5

·

1 Parent(s): 07a35ff

Fix code example

Files changed (1) hide show

README.md +2 -1

README.md CHANGED Viewed

@@ -38,7 +38,8 @@ url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
 image = Image.open(requests.get(url, stream=True).raw)
 feature_extractor = AutoFeatureExtractor.from_pretrained('facebook/vit-mae-base')
-model = ViTModel.from_pretrained('facebook/vit-mae-base')
 inputs = feature_extractor(images=image, return_tensors="pt")
 outputs = model(**inputs)
 loss = outputs.loss

 image = Image.open(requests.get(url, stream=True).raw)
 feature_extractor = AutoFeatureExtractor.from_pretrained('facebook/vit-mae-base')
+model = ViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base')
 inputs = feature_extractor(images=image, return_tensors="pt")
 outputs = model(**inputs)
 loss = outputs.loss