Update model config and README

README.md

@@ -19,9 +19,9 @@ An EfficientViT (MIT) image classification model. Trained on ImageNet-1k by pape
   - Activations (M): 9.5
   - Image size: 256 x 256
 - **Papers:**
-  - EfficientViT: Lightweight Multi-Scale Attention for On-Device Semantic Segmentation: https://arxiv.org/abs/2205.14756
-- **Dataset:** ImageNet-1k
+  - EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction: https://arxiv.org/abs/2205.14756
 - **Original:** https://github.com/mit-han-lab/efficientvit
+- **Dataset:** ImageNet-1k
 
 ## Model Usage
 ### Image Classification
@@ -112,13 +112,10 @@ output = model.forward_head(output, pre_logits=True)
 # output is a (1, num_features) shaped tensor
 ```
 
-## Model Comparison
-Explore the dataset and runtime metrics of this model in timm [model results](https://github.com/huggingface/pytorch-image-models/tree/main/results).
-
 ## Citation
 ```bibtex
 @article{cai2022efficientvit,
-  title={Efficientvit: Enhanced linear attention for high-resolution low-computation visual recognition},
+  title={EfficientViT: Enhanced linear attention for high-resolution low-computation visual recognition},
   author={Cai, Han and Gan, Chuang and Han, Song},
   journal={arXiv preprint arXiv:2205.14756},
   year={2022}

config.json

@@ -26,7 +26,10 @@
       0.225
     ],
     "num_classes": 1000,
-    "pool_size": null,
+    "pool_size": [
+      8,
+      8
+    ],
     "first_conv": "stem.in_conv.conv",
     "classifier": "head.classifier.4"
   }