Apply suggestions from code review

huggingface · Dec 2, 2024 · c500669 · c500669
1 parent 3ea0329
commit c500669
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -432,8 +432,8 @@ class TasksManager:
             tflite="BertTFLiteConfig",
         ),
         "rembert": supported_tasks_mapping(
-            "text-generation",
             "fill-mask",
+            "feature-extraction",
             "text-classification",
             "multiple-choice",
             "token-classification",

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -2331,7 +2331,6 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin):
         "llama",
         "mistral",
         "opt",
-        "rembert",
     ]
 
     if check_if_transformers_greater("4.37"):