Fix test_load_default_pipelines_tf test error (huggingface#18422)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
oneraghavan · Sep 26, 2022 · 42cb243 · 42cb243
1 parent 341b0af
commit 42cb243
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py
@@ -285,7 +285,12 @@
         "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": {