nit

huggingface · Aug 23, 2023 · f2bc80e · f2bc80e
1 parent 878a0b7
commit f2bc80e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/exporters/onnx/test_onnx_export.py b/tests/exporters/onnx/test_onnx_export.py
@@ -338,7 +338,8 @@ def _onnx_export_sd(self, model_type: str, model_name: str, device="cpu"):
     def test_all_models_tested(self):
         # make sure we test all models
         missing_models_set = TasksManager._SUPPORTED_CLI_MODEL_TYPE - set(PYTORCH_EXPORT_MODELS_TINY.keys())
-        if len(missing_models_set) > 0:
+        assert "sam" in missing_models_set  # See exporters_utils.py
+        if len(missing_models_set) > 1:
             self.fail(f"Not testing all models. Missing models: {missing_models_set}")
 
     @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS_TINY))