ybelkada commited on
Commit
60f75d5
1 Parent(s): c31a503

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +5 -4
README.md CHANGED
@@ -32,15 +32,15 @@ fine-tuned versions on a task that interests you.
32
  Here is how to use this model to classify an image of the COCO 2017 dataset into one of the 1,000 ImageNet classes:
33
 
34
  ```python
35
- from transformers import AutoFeatureExtractor, BiTForImageClassification
36
  import torch
37
  from datasets import load_dataset
38
 
39
  dataset = load_dataset("huggingface/cats-image")
40
  image = dataset["test"]["image"][0]
41
 
42
- feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
43
- model = BiTForImageClassification.from_pretrained("google/bit-50")
44
 
45
  inputs = feature_extractor(image, return_tensors="pt")
46
 
@@ -49,7 +49,8 @@ with torch.no_grad():
49
 
50
  # model predicts one of the 1000 ImageNet classes
51
  predicted_label = logits.argmax(-1).item()
52
- print(model.config.id2label[predicted_label])
 
53
  ```
54
 
55
  For more code examples, we refer to the [documentation](https://huggingface.co/docs/transformers/main/en/model_doc/bit).
 
32
  Here is how to use this model to classify an image of the COCO 2017 dataset into one of the 1,000 ImageNet classes:
33
 
34
  ```python
35
+ from transformers import BitImageProcessor, BitForImageClassification
36
  import torch
37
  from datasets import load_dataset
38
 
39
  dataset = load_dataset("huggingface/cats-image")
40
  image = dataset["test"]["image"][0]
41
 
42
+ feature_extractor = BitImageProcessor.from_pretrained("google/bit-50")
43
+ model = BitForImageClassification.from_pretrained("google/bit-50")
44
 
45
  inputs = feature_extractor(image, return_tensors="pt")
46
 
 
49
 
50
  # model predicts one of the 1000 ImageNet classes
51
  predicted_label = logits.argmax(-1).item()
52
+ print(model.config.id2label[predicted_label
53
+ >>> tabby, tabby cat
54
  ```
55
 
56
  For more code examples, we refer to the [documentation](https://huggingface.co/docs/transformers/main/en/model_doc/bit).