Add back FE pipeline for Vilt

huggingface · Jan 31, 2024 · f9daf11 · f9daf11
1 parent c7a3e92
commit f9daf11
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/models/vilt/test_modeling_vilt.py b/tests/models/vilt/test_modeling_vilt.py
@@ -227,7 +227,7 @@ class ViltModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
         if is_torch_available()
         else ()
     )
-    pipeline_model_mapping = {"visual-question-answering": ViltForQuestionAnswering} if is_torch_available() else {}
+    pipeline_model_mapping = {"image-feature-extraction": ViltModel, "visual-question-answering": ViltForQuestionAnswering} if is_torch_available() else {}
     test_pruning = False
     test_headmasking = False
     test_torchscript = False