Fix `test_load_default_pipelines_tf` test error (#18422)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Fix `test_load_default_pipelines_tf` test error (#18422)
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
2959d090 · Yih-Dar · GitHub · 8ae77842 · 2959d090
Unverified Commit 2959d090 authored Aug 02, 2022 by Yih-Dar Committed by GitHub Aug 02, 2022
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

src/transformers/pipelines/__init__.py src/transformers/pipelines/__init__.py +6 -1

No files found.
--- a/src/transformers/pipelines/__init__.py
+++ b/src/transformers/pipelines/__init__.py
@@ -285,7 +285,12 @@ SUPPORTED_TASKS = {
        "impl": ImageClassificationPipeline,
        "tf": (TFAutoModelForImageClassification,) if is_tf_available() else (),
        "pt": (AutoModelForImageClassification,) if is_torch_available() else (),
-        "default": {"model": {"pt": ("google/vit-base-patch16-224", "5dca96d")}},
+        "default": {
+            "model": {
+                "pt": ("google/vit-base-patch16-224", "5dca96d"),
+                "tf": ("google/vit-base-patch16-224", "5dca96d"),
+            }
+        },
        "type": "image",
    },
    "image-segmentation": {