nielsr HF staff commited on
Commit
54a0cf5
1 Parent(s): 07a35ff

Fix code example

Browse files
Files changed (1) hide show
  1. README.md +2 -1
README.md CHANGED
@@ -38,7 +38,8 @@ url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
38
  image = Image.open(requests.get(url, stream=True).raw)
39
 
40
  feature_extractor = AutoFeatureExtractor.from_pretrained('facebook/vit-mae-base')
41
- model = ViTModel.from_pretrained('facebook/vit-mae-base')
 
42
  inputs = feature_extractor(images=image, return_tensors="pt")
43
  outputs = model(**inputs)
44
  loss = outputs.loss
 
38
  image = Image.open(requests.get(url, stream=True).raw)
39
 
40
  feature_extractor = AutoFeatureExtractor.from_pretrained('facebook/vit-mae-base')
41
+ model = ViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base')
42
+
43
  inputs = feature_extractor(images=image, return_tensors="pt")
44
  outputs = model(**inputs)
45
  loss = outputs.loss