oh no!!!

huggingface · Jan 8, 2025 · 1b42ba1 · 1b42ba1
1 parent de7664d
commit 1b42ba1
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py
@@ -1410,10 +1410,12 @@ def assert_screenout(out, what):
 
 def set_model_tester_for_less_flaky_test(test_case):
     if hasattr(test_case.model_tester, "num_hidden_layers"):
-        test_case.model_tester.num_hidden_layers = 1
+        if not (hasattr(test_case.model_tester, "out_features") or hasattr(test_case.model_tester, "out_indices")):
+            test_case.model_tester.num_hidden_layers = 1
     if (
         hasattr(test_case.model_tester, "vision_config")
         and "num_hidden_layers" in test_case.model_tester.vision_config
+        and (not (hasattr(test_case.model_tester.vision_config, "out_features") or hasattr(test_case.model_tester.vision_config, "out_indices")))
     ):
         test_case.model_tester.vision_config = copy.deepcopy(test_case.model_tester.vision_config)
         test_case.model_tester.vision_config["num_hidden_layers"] = 1
@@ -1431,12 +1433,6 @@ def set_model_tester_for_less_flaky_test(test_case):
         test_case.model_tester.neck_hidden_sizes = test_case.model_tester.neck_hidden_sizes[
             : test_case.model_tester.num_hidden_layers
         ]
-    # For DPT auto backbone
-    if hasattr(test_case.model_tester, "num_hidden_layers"):
-        if hasattr(test_case.model_tester, "out_features"):
-            test_case.model_tester.out_features = test_case.model_tester.out_features[
-                : test_case.model_tester.num_hidden_layers
-            ]
 
 
 def set_config_for_less_flaky_test(config):