recompile limit

Author: sayakpaul

tests/models/testing_utils/compile.py

@@ -92,9 +92,6 @@ def test_torch_compile_repeated_blocks(self, recompile_limit=1):
         model.eval()
         model.compile_repeated_blocks(fullgraph=True)
 
-        if self.model_class.__name__ == "UNet2DConditionModel":
-            recompile_limit = 2
-
         with (
             torch._inductor.utils.fresh_inductor_cache(),
             torch._dynamo.config.patch(recompile_limit=recompile_limit),

tests/models/unets/test_models_unet_2d_condition.py

@@ -1199,6 +1199,9 @@ def test_ip_adapter_plus(self):
 class TestUNet2DConditionModelCompile(UNet2DConditionTesterConfig, TorchCompileTesterMixin):
     """Torch compile tests for UNet2DConditionModel."""
 
+    def test_torch_compile_repeated_blocks(self):
+        return super().test_torch_compile_repeated_blocks(recompile_limit=2)
+
 
 class TestUNet2DConditionModelLoRAHotSwap(UNet2DConditionTesterConfig, LoraHotSwappingForModelTesterMixin):
     """LoRA hot-swapping tests for UNet2DConditionModel."""