Add missing YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP

huggingface · Apr 20, 2022 · f414e10 · f414e10
1 parent 960d324
commit f414e10
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/transformers/models/yolos/configuration_yolos.py b/src/transformers/models/yolos/configuration_yolos.py
@@ -20,9 +20,9 @@
 
 logger = logging.get_logger(__name__)
 
-VIT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "google/vit-base-patch16-224": "https://huggingface.co/vit-base-patch16-224/resolve/main/config.json",
-    # See all YOLOS models at https://huggingface.co/models?filter=vit
+YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "google/yolos-s": "https://huggingface.co/yolos-s/resolve/main/config.json",
+    # See all YOLOS models at https://huggingface.co/models?filter=yolos
 }
 
 

diff --git a/src/transformers/models/yolos/modeling_yolos.py b/src/transformers/models/yolos/modeling_yolos.py
@@ -544,7 +544,7 @@ class YolosPreTrainedModel(PreTrainedModel):
     """
 
     config_class = YolosConfig
-    base_model_prefix = "yolos"
+    base_model_prefix = "vit"
     main_input_name = "pixel_values"
     supports_gradient_checkpointing = True