[vision] Add problem_type support (#15851)

* Add problem_type to missing models

* Fix deit test

Co-authored-by: Niels Rogge <nielsrogge@Nielss-MacBook-Pro.local>
This commit is contained in:
NielsRogge
2022-03-01 18:09:52 +01:00
committed by GitHub
parent 7ff9d450cd
commit 286fdc6b3c
6 changed files with 130 additions and 25 deletions

View File

@@ -1873,7 +1873,10 @@ class ModelTesterMixin:
]
for model_class in self.all_model_classes:
if model_class not in get_values(MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING):
if model_class not in [
*get_values(MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING),
*get_values(MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING),
]:
continue
for problem_type in problem_types: